# Regression

## Auxiliary

function lreg_chkdims(X::AbstractMatrix, Y::AbstractVecOrMat, trans::Bool)
    mX, nX = size(X)
    dX = ifelse(trans, mX, nX)
    dY = ifelse(trans, nX, mX)
    size(Y, 1) == dY || throw(DimensionMismatch("Dimensions of X and Y mismatch."))
    return dX
end

lrsoltype(::AbstractVector{T}) where T = Vector{T}
lrsoltype(::AbstractMatrix{T}) where T = Matrix{T}

_vaug(X::AbstractMatrix{T}) where T = vcat(X, ones(T, 1, size(X,2)))::Matrix{T}
_haug(X::AbstractMatrix{T}) where T = hcat(X, ones(T, size(X,1), 1))::Matrix{T}

## Linear Least Square Regression

"""
    llsq(X, y; ...)

Solve the linear least square problem.

Here, `y` can be either a vector, or a matrix where each column is a response vector.

This function accepts two keyword arguments:

- `dims`: whether input observations are stored as rows (`1`) or columns (`2`). (default is `1`)
- `bias`: whether to include the bias term `b`. (default is `true`)

The function results the solution `a`. In particular, when `y` is a vector (matrix), `a` is also a vector (matrix). If `bias` is true, then the returned array is augmented as `[a; b]`.
"""
function llsq(X::AbstractMatrix{T}, Y::AbstractVecOrMat{T};
              trans::Bool=false, bias::Bool=true,
              dims::Union{Integer,Nothing}=nothing) where {T<:Real}
    if dims === nothing && trans
        Base.depwarn("`trans` argument is deprecated, use llsq(X, Y, dims=d) instead.", :trans)
        dims = 1
    end
    if dims == 2
        mX, nX = size(X)
        size(Y, 1) == nX || throw(DimensionMismatch("Dimensions of X and Y mismatch."))
        mX <= nX || error("mX <= nX is required when trans is false.")
    else
        mX, nX = size(X)
        size(Y, 1) == mX || throw(DimensionMismatch("Dimensions of X and Y mismatch."))
        mX >= nX || error("mX >= nX is required when trans is false.")
    end
    _ridge(X, Y, zero(T), dims == 2, bias)
end
llsq(x::AbstractVector{T}, y::AbstractVector{T}; kwargs...) where {T<:Real} =
    llsq(x[:,:], y; dims=1, kwargs...)

## Ridge Regression (Tikhonov regularization)

"""

    ridge(X, y, r; ...)

Solve the ridge regression problem.

Here, ``y`` can be either a vector, or a matrix where each column is a response vector.

The argument `r` gives the quadratic regularization matrix ``Q``, which can be in either of the following forms:

- `r` is a real scalar, then ``Q`` is considered to be `r * eye(n)`, where `n` is the dimension of `a`.
- `r` is a real vector, then ``Q`` is considered to be `diagm(r)`.
- `r` is a real symmetric matrix, then ``Q`` is simply considered to be `r`.

This function accepts two keyword arguments:

- `dims`: whether input observations are stored as rows (`1`) or columns (`2`). (default is `1`)
- `bias`: whether to include the bias term `b`. (default is `true`)

The function results the solution `a`. In particular, when `y` is a vector (matrix), `a` is also a vector (matrix). If `bias` is true, then the returned array is augmented as `[a; b]`.
"""
function ridge(X::AbstractMatrix{T}, Y::AbstractVecOrMat{T}, r::Union{Real, AbstractVecOrMat};
               trans::Bool=false, bias::Bool=true,
               dims::Union{Integer,Nothing}=nothing) where {T<:Real}
    if dims === nothing && trans
        Base.depwarn("`trans` argument is deprecated, use ridge(X, Y, r, dims=d) instead.", :trans)
        dims = 1
    end
    d = lreg_chkdims(X, Y, dims == 2)
    if isa(r, Real)
        r >= zero(r) || error("r must be non-negative.")
        r = convert(T <: AbstractFloat ? T : Float64, r)
    elseif isa(r, AbstractVector)
        length(r) == d || throw(DimensionMismatch("Incorrect length of r."))
    elseif isa(r, AbstractMatrix)
        size(r) == (d, d) || throw(DimensionMismatch("Incorrect size of r."))
    end
    _ridge(X, Y, r, dims == 2, bias)
end
ridge(x::AbstractVector{T}, y::AbstractVector{T}, r::Union{Real, AbstractVecOrMat};
      kwargs...) where {T<:Real} = ridge(x[:,:], y, r; dims=1, kwargs...)

### implementation

function _ridge(_X::AbstractMatrix{T}, _Y::AbstractVecOrMat{T},
                r::Union{Real, AbstractVecOrMat}, trans::Bool, bias::Bool) where {T<:Real}
    # convert integer data to Float64
    X = T <: AbstractFloat ? _X : convert(Array{Float64}, _X)
    Y = T <: AbstractFloat ? _Y : convert(Array{Float64}, _Y)
    if bias
        if trans
            X_ = _vaug(X)
            A = cholesky!(Hermitian(_ridge_reg!(X_ * transpose(X_), r, bias))) \ (X_ * Y)
        else
            X_ = _haug(X)
            A = cholesky!(Hermitian(_ridge_reg!(X_'X_, r, bias))) \ (X_'Y)
        end
    else
        if trans
            A = cholesky!(Hermitian(_ridge_reg!(X * X', r, bias))) \ (X * Y)
        else
            A = cholesky!(Hermitian(_ridge_reg!(X'X, r, bias))) \ (X'Y)
        end
    end
    return A::lrsoltype(Y)
end

function _ridge_reg!(Q::Matrix, r::Real, bias::Bool)
    if r > zero(r)
        n = size(Q, 1) - Int(bias)
        for i = 1:n
            @inbounds Q[i,i] += r
        end
    end
    return Q
end

function _ridge_reg!(Q::AbstractMatrix, r::AbstractVector, bias::Bool)
    n = size(Q, 1) - Int(bias)
    @assert length(r) == n
    for i = 1:n
        @inbounds Q[i,i] += r[i]
    end
    return Q
end

function _ridge_reg!(Q::AbstractMatrix, r::AbstractMatrix, bias::Bool)
    n = size(Q, 1) - Int(bias)
    @assert size(r) == (n, n)
    for j = 1:n, i = 1:n
        @inbounds Q[i,j] += r[i,j]
    end
    return Q
end

## Isotonic Regression

"""
    isotonic(x, y[, w])

Solve the isotonic regression problem using the pool adjacent violators algorithm[^1].

Here `x` is the regressor vector, `y` is response vector, and `w` is an optional
weights vector.

The function returns a prediction vector of the same size as the regressor vector `x`.
"""
function isotonic(x::AbstractVector{T}, y::AbstractVector{T},
                  w::AbstractVector{T} = ones(T, length(y))) where {T<:Real}
    n = length(x)
    n == length(y) || throw(DimensionMismatch("Dimensions of x and y mismatch."))

    idx = sortperm(x)

    # PVA algorithm
    J = map(i->(Float64(y[i]*w[i]), w[i], [i]), idx)
    i = 1
    B₀ = J[i]
    while i < length(J)
        B₊ = J[i+1]
        if B₀[1] <= B₊[1] # step 1
            B₀ = B₊
            i += 1
        else # step 2
            ww = B₀[2] + B₊[2]
            J[i] = ((B₀[1]*B₀[2]+B₊[1]*B₊[2])/ww, ww, append!(B₀[3], B₊[3]))
            deleteat!(J, i+1)
            B₀ = J[i]
            while i > 1 # step 2.1
                B₋ = J[i-1]
                if B₀[1] <= B₋[1]
                    ww = B₀[2] + B₋[2]
                    J[i] = ((B₀[1]*B₀[2]+B₋[1]*B₋[2])/ww, ww, append!(B₀[3], B₋[3]))
                    deleteat!(J, i-1)
                    i-=1
                else
                    break
                end
            end
        end
    end
    [y for (y,w,ii) in J for i in sort(ii)]
end


#=
Created on 09/01/2021 14:50:19
Last update: -

@author: Michiel Stock
michielfmstock@gmail.com

Implement Stirling's approximation for the factorial.
=#

using Plots

stirling(n) = n * log(n / exp(1)) + 0.5log(2π*n)

# here we use `m=:circle` to use the :circle symbols as a marker on our plots.

plot(1:10, n->log(factorial(big(n))), m=:circle, label="log-factorial")
xlabel!("n")
plot!(stirling, 1:10, m=:circle, label = "Stirling's approximation")


#   Unit classification.jl, part of PosDefManifold Package for julia language
#   v 0.1.3 - last update 7th of Mai 2019
#
#   MIT License
#   Copyright (c) 2019, Marco Congedo, CNRS, Grenobe, France:
#   https://sites.google.com/site/marcocongedo/home
#
#   DESCRIPTION
#   This Unit implements classification methods and related functions
#   useful in relation to the Riemannian Geometry on the manifold
#   of Symmetric Positive Definite (SPD) or Hermitian matrices.
#
#   CONTENT
#   1. Probability
# __________________________________________________________________

#  ------------------------
## 1. Probability
#  ------------------------

"""
    softmax(χ::Vector{T}) where T<:Real

 Given a real vector of ``k`` non-negative scores ``χ=c_1,...,c_k``,
 return the vector ``π=p_1,...,p_k`` of their
 [softmax](https://en.wikipedia.org/wiki/Softmax_function) probabilities,
 as per

 `` p_i=\\frac{\\textrm{e}^{c_i}}{\\sum_{i=1}^{k}\\textrm{e}^{c_i}} ``.

 ## Examples
    χ=[1.0, 2.3, 0.4, 5.0]
    π=softmax(χ)

"""
softmax(χ::Vector{T}) where T<:Real = exp.(χ) ./ 𝚺(exp.(χ))


export hermite, HermiteBasis

"""
    y = hermite(n::Int, x::Real)

Evaluate the ``n``-degree Hermite polynomial at ``x``.

```math
n!\\sum_{m=0}^{[n/2]} \\frac{(-1)^m}{m!(n-2m)!} \\frac{x^{n-2m}}{2^m}
```
"""
function hermite(n::Int, x::Real)
  if n == 0
    1
  else
    N = Int(floor(n/2))
    fac = factorial
    s = reduce((s, m) -> s + (-1)^m * x^(n-2m) / (fac(m) * fac(n-2m) * 2^m), 0:N; init=0)
    fac(n) * s
  end
end

"""
    basis = HermiteBasis(vectors)

Encode an ordered basis (`basis`) of polynomials ``(ψ_j)_{j\\in \\mathcal J_i}`` of the form

```math
ψ_{(j_1,\\ldots,j_n)}(θ_1, \\ldots, θ_n) = \\prod_{k=1}^n φ_{j_k}(θ_k),
```

where ``φ_{j_k}`` is the Hermite polynomial of degree ``j_k``.
The indices ``\\mathcal J_i`` of the polynomials is provided by `vectors` in the natural
way. For example,

```julia
HermiteBasis([[1, 2, 3], [0, 1, 1, 2]])
```

will encode the ordered basis ``( ψ_{(1,2,3)}, ψ_{(0,1,1,2)} )``.
"""
struct HermiteBasis <: OrderedBasis
  vectors::AbstractArray{<:AbstractArray{<:Int}}
  function HermiteBasis(vectors::AbstractArray{<:AbstractArray{<:Int}})
    if all(vec -> all(x -> x >= 0, vec), vectors)
      new(vectors)
    else
      throw(
        ArgumentError("A `HermiteBasis` must only include vectors with positive components.")
       )
    end
  end
end

Base.length(basis::HermiteBasis) = length(basis.vectors)

function Base.union(basis::HermiteBasis, bases...)
  HermiteBasis(union(basis.vectors, map(b -> b.vectors, bases)...))
end

function (basis::HermiteBasis)(theta::AbstractArray{<:Real})
  if any(vec -> length(vec) > length(theta), basis.vectors)
    throw(ArgumentError("$theta is not in the domain of $basis"))
  else
    map(vec -> prod(map(p -> hermite(p[2], theta[p[1]]), enumerate(vec))), basis.vectors)
  end
end


# # Gaussian process prior samples
#
# The kernels defined in this package can also be used to specify the
# covariance of a Gaussian process prior.
# A Gaussian process (GP) is defined by its mean function $m(\cdot)$ and its covariance function or kernel $k(\cdot, \cdot')$:
# ```math
#   f \sim \mathcal{GP}\big(m(\cdot), k(\cdot, \cdot')\big)
# ```
# In this notebook we show how the choice of kernel affects the samples from a GP (with zero mean).

## Load required packages
using KernelFunctions, LinearAlgebra
using Plots, Plots.PlotMeasures
default(; lw=1.0, legendfontsize=8.0)
using Random: seed!
seed!(42); # reproducibility

##
# ## Evaluation at finite set of points
#
# The function values $\mathbf{f} = \{f(x_n)\}_{n=1}^N$ of the GP at a finite number $N$ of points $X = \{x_n\}_{n=1}^N$ follow a multivariate normal distribution $\mathbf{f} \sim \mathcal{MVN}(\mathbf{m}, \mathrm{K})$ with mean vector $\mathbf{m}$ and covariance matrix $\mathrm{K}$, where
# ```math
# \begin{aligned}
#   \mathbf{m}_i &= m(x_i) \\
#   \mathrm{K}_{i,j} &= k(x_i, x_j)
# \end{aligned}
# ```
# with $1 \le i, j \le N$.
#
# We can visualize the infinite-dimensional GP by evaluating it on a fine grid to approximate the dense real line:
num_inputs = 101
xlim = (-5, 5)
X = range(xlim...; length=num_inputs);

# Given a kernel `k`, we can compute the kernel matrix as `K = kernelmatrix(k, X)`.

##
# ## Random samples
#
# To sample from the multivariate normal distribution $p(\mathbf{f}) = \mathcal{MVN}(0, \mathrm{K})$, we could make use of Distributions.jl and call `rand(MvNormal(K))`.
# Alternatively, we could use the [AbstractGPs.jl](https://github.com/JuliaGaussianProcesses/AbstractGPs.jl) package and construct a `GP` object which we evaluate at the points of interest and from which we can then sample: `rand(GP(k)(X))`.

# Here, we will explicitly construct samples using the Cholesky factorization $\mathrm{L} = \operatorname{cholesky}(\mathrm{K})$,
# with $\mathbf{f} = \mathrm{L} \mathbf{v}$, where $\mathbf{v} \sim \mathcal{N}(0, \mathbf{I})$ is a vector of standard-normal random variables.

# We will use the same randomness $\mathbf{v}$ to generate comparable samples across different kernels.
num_samples = 7
v = randn(num_inputs, num_samples);

# Mathematically, a kernel matrix is by definition positive semi-definite, but due to finite-precision inaccuracies, the computed kernel matrix might not be exactly positive definite. To avoid Cholesky errors, we add a small "nugget" term on the diagonal:
function mvn_sample(K)
    L = cholesky(K + 1e-6 * I)
    f = L.L * v
    return f
end;

##
# ## Visualization
# We now define a function that visualizes a kernel for us.

function visualize(k::Kernel)
    K = kernelmatrix(k, X)
    f = mvn_sample(K)

    p_kernel_2d = heatmap(
        X,
        X,
        K;
        yflip=true,
        colorbar=false,
        ylabel=string(nameof(typeof(k))),
        ylim=xlim,
        yticks=([xlim[1], 0, xlim[end]], ["\u22125", raw"$x'$", "5"]),
        vlim=(0, 1),
        title=raw"$k(x, x')$",
        aspect_ratio=:equal,
        left_margin=5mm,
    )

    p_kernel_cut = plot(
        X,
        k.(X, 0.0);
        title=string(raw"$k(x, x_\mathrm{ref})$"),
        label=raw"$x_\mathrm{ref}=0.0$",
        legend=:topleft,
        foreground_color_legend=nothing,
    )
    plot!(X, k.(X, 1.5); label=raw"$x_\mathrm{ref}=1.5$")

    p_samples = plot(X, f; c="blue", title=raw"$f(x)$", ylim=(-3, 3), label=nothing)

    return plot(
        p_kernel_2d,
        p_kernel_cut,
        p_samples;
        layout=(1, 3),
        xlabel=raw"$x$",
        xlim=xlim,
        xticks=collect(xlim),
    )
end;

##
# We can now visualize a kernel and show samples from
# a Gaussian process with a given kernel:

plot(visualize(SqExponentialKernel()); size=(800, 210), bottommargin=5mm, topmargin=5mm)

##
# ## Kernel comparison
# This also allows us to compare different kernels:

kernels = [
    Matern12Kernel(),
    Matern32Kernel(),
    Matern52Kernel(),
    SqExponentialKernel(),
    WhiteKernel(),
    ConstantKernel(),
    LinearKernel(),
    compose(PeriodicKernel(), ScaleTransform(0.2)),
    NeuralNetworkKernel(),
]
plot(
    [visualize(k) for k in kernels]...;
    layout=(length(kernels), 1),
    size=(800, 220 * length(kernels) + 100),
)


module NewtonRaphson

export newton_raphson

function newton_raphson(f::Function, x0::Number, fprime::Function, args::Tuple=();
                tol::AbstractFloat=1e-8, maxiter::Integer=50, eps::AbstractFloat=1e-10)
    for _ in 1:maxiter
        yprime = fprime(x0, args...)
        if abs(yprime) < eps
           warn("First derivative is zero")
           return x0
        end
        y = f(x0, args...)
        x1 = x0 - y/yprime
        if abs(x1-x0) < tol
          return x1
        end
        x0 = x1
    end
    error("Max iteration exceeded")
end

end


# # An Intro to DifferentialEquations.jl
#
# ## Basic Introduction Via Ordinary Differential Equations
#
# This notebook will get you started with DifferentialEquations.jl by introducing you to the functionality for solving ordinary differential equations (ODEs). The corresponding documentation page is the [ODE tutorial](http://docs.juliadiffeq.org/latest/tutorials/ode_example.html). While some of the syntax may be different for other types of equations, the same general principles hold in each case. Our goal is to give a gentle and thorough introduction that highlights these principles in a way that will help you generalize what you have learned.
#
# ### Background
#
# If you are new to the study of differential equations, it can be helpful to do a quick background read on [the definition of ordinary differential equations](https://en.wikipedia.org/wiki/Ordinary_differential_equation). We define an ordinary differential equation as an equation which describes the way that a variable $u$ changes, that is
#
# $$ u' = f(u,p,t) $$
#
# where $p$ are the parameters of the model, $t$ is the time variable, and $f$ is the nonlinear model of how $u$ changes. The initial value problem also includes the information about the starting value:
#
# $$ u(t_0) = u_0 $$
#
# Together, if you know the starting value and you know how the value will change with time, then you know what the value will be at any time point in the future. This is the intuitive definition of a differential equation.

# ### First Model: Exponential Growth
#
# Our first model will be the canonical exponential growth model. This model says that the rate of change is proportional to the current value, and is this:
#
# $$ u' = au $$
#
# where we have a starting value $u(0)=u_0$. Let's say we put 1 dollar into Bitcoin which is increasing at a rate of $98\%$ per year. Then calling now $t=0$ and measuring time in years, our model is:
#
# $$ u' = 0.98u $$
#
# and $u(0) = 1.0$. We encode this into Julia by noticing that, in this setup, we match the general form when
#
# $$ f(u,p,t) = 0.98u $$
#
# with $ u_0 = 1.0 $. If we want to solve this model on a time span from `t=0.0` to `t=1.0`, then we define an `ODEProblem` by specifying this function `f`, this initial condition `u0`, and this time span as follows:

using DifferentialEquations
f(u,p,t) = 0.98u
u0 = 1.0
tspan = (0.0,1.0)
prob = ODEProblem(f,u0,tspan)
#----------------------------------------------------------------------------

# To solve our `ODEProblem` we use the command `solve`.

sol = solve(prob)
#----------------------------------------------------------------------------

# and that's it: we have succesfully solved our first ODE!
#
# #### Analyzing the Solution
#
# Of course, the solution type is not interesting in and of itself. We want to understand the solution! The documentation page which explains in detail the functions for analyzing the solution is the [Solution Handling](http://docs.juliadiffeq.org/latest/basics/solution.html) page. Here we will describe some of the basics. You can plot the solution using the plot recipe provided by [Plots.jl](http://docs.juliaplots.org/latest/):

using Plots; gr()
plot(sol)
#----------------------------------------------------------------------------

# From the picture we see that the solution is an exponential curve, which matches our intuition. As a plot recipe, we can annotate the result using any of the [Plots.jl attributes](http://docs.juliaplots.org/latest/attributes/). For example:

plot(sol,linewidth=5,title="Solution to the linear ODE with a thick line",
     xaxis="Time (t)",yaxis="u(t) (in μm)",label="My Thick Line!") ## legend=false
#----------------------------------------------------------------------------

# Using the mutating `plot!` command we can add other pieces to our plot. For this ODE we know that the true solution is $u(t) = u_0 exp(at)$, so let's add some of the true solution to our plot:

plot!(sol.t, t->1.0*exp(0.98t),lw=3,ls=:dash,label="True Solution!")
#----------------------------------------------------------------------------

# In the previous command I demonstrated `sol.t`, which grabs the array of time points that the solution was saved at:

sol.t
#----------------------------------------------------------------------------

# We can get the array of solution values using `sol.u`:

sol.u
#----------------------------------------------------------------------------

# `sol.u[i]` is the value of the solution at time `sol.t[i]`. We can compute arrays of functions of the solution values using standard comprehensions, like:

[t+u for (u,t) in tuples(sol)]
#----------------------------------------------------------------------------

# However, one interesting feature is that, by default, the solution is a continuous function. If we check the print out again:

sol
#----------------------------------------------------------------------------

# you see that it says that the solution has a 4th order interpolation, meaning that it is a continuous function of 4th order accuracy. We can call the solution as a function of time `sol(t)`. For example, to get the value at `t=0.45`, we can use the command:

sol(0.45)
#----------------------------------------------------------------------------

# #### Controlling the Solver
#
# DifferentialEquations.jl has a common set of solver controls among its algorithms which can be found [at the Common Solver Options](http://docs.juliadiffeq.org/latest/basics/common_solver_opts.html) page. We will detail some of the most widely used options.
#
# The most useful options are the tolerances `abstol` and `reltol`. These tell the internal adaptive time stepping engine how precise of a solution you want. Generally, `reltol` is the relative accuracy while `abstol` is the accuracy when `u` is near zero. These tolerances are local tolerances and thus are not global guarantees. However, a good rule of thumb is that the total solution accuracy is 1-2 digits less than the relative tolerances. Thus for the defaults `abstol=1e-6` and `reltol=1e-3`, you can expect a global accuracy of about 1-2 digits. If we want to get around 6 digits of accuracy, we can use the commands:

sol = solve(prob,abstol=1e-8,reltol=1e-8)
#----------------------------------------------------------------------------

# Now we can see no visible difference against the true solution:

plot(sol)
plot!(sol.t, t->1.0*exp(0.98t),lw=3,ls=:dash,label="True Solution!")
#----------------------------------------------------------------------------

# Notice that by decreasing the tolerance, the number of steps the solver had to take was `9` instead of the previous `5`. There is a trade off between accuracy and speed, and it is up to you to determine what is the right balance for your problem.
#
# Another common option is to use `saveat` to make the solver save at specific time points. For example, if we want the solution at an even grid of `t=0.1k` for integers `k`, we would use the command:

sol = solve(prob,saveat=0.1)
#----------------------------------------------------------------------------

# Notice that when `saveat` is used the continuous output variables are no longer saved and thus `sol(t)`, the interpolation, is only first order. We can save at an uneven grid of points by passing a collection of values to `saveat`. For example:

sol = solve(prob,saveat=[0.2,0.7,0.9])
#----------------------------------------------------------------------------

# By default it always saves the first and last values, but we can turn this off as well:

sol = solve(prob,saveat=[0.2,0.7,0.9],save_start = false, save_end = false)
#----------------------------------------------------------------------------

# If we need to reduce the amount of saving, we can also turn off the continuous output directly via `dense=false`:

sol = solve(prob,dense=false)
#----------------------------------------------------------------------------

# and to turn off all intermediate saving we can use `save_everystep=false`:

sol = solve(prob,save_everystep=false)
#----------------------------------------------------------------------------

# More advanced saving behaviors, such as saving functionals of the solution, are handled via the `SavingCallback` in the [Callback Library](http://docs.juliadiffeq.org/latest/features/callback_library.html#SavingCallback-1) which will be addressed later in the tutorial.

# #### Choosing Solver Algorithms
#
# There is no best algorithm for numerically solving a differential equation. When you call `solve(prob)`, DifferentialEquations.jl makes a guess at a good algorithm for your problem, given the properties that you ask for (the tolerances, the saving information, etc.). However, in many cases you may want more direct control. A later notebook will help introduce the various *algorithms* in DifferentialEquations.jl, but for now let's introduce the *syntax*.
#
# The most crucial determining factor in choosing a numerical method is the stiffness of the model. Stiffness is roughly characterized by a Jacobian `f` with large eigenvalues. That's quite mathematical, and we can think of it more intuitively: if you have big numbers in `f` (like parameters of order `1e5`), then it's probably stiff. Or, as the creator of the MATLAB ODE Suite, Lawrence Shampine, likes to define it, if the standard algorithms are slow, then it's stiff. We will go into more depth about diagnosing stiffness in a later tutorial, but for now note that if you believe your model may be stiff, you can hint this to the algorithm chooser via `alg_hints = [:stiff]`.

sol = solve(prob,alg_hints=[:stiff])
#----------------------------------------------------------------------------

# Stiff algorithms have to solve implicit equations and linear systems at each step so they should only be used when required.
#
# If we want to choose an algorithm directly, you can pass the algorithm type after the problem as `solve(prob,alg)`. For example, let's solve this problem using the `Tsit5()` algorithm, and just for show let's change the relative tolerance to `1e-6` at the same time:

sol = solve(prob,Tsit5(),reltol=1e-6)
#----------------------------------------------------------------------------

# ### Systems of ODEs: The Lorenz Equation
#
# Now let's move to a system of ODEs. The [Lorenz equation](https://en.wikipedia.org/wiki/Lorenz_system) is the famous "butterfly attractor" that spawned chaos theory. It is defined by the system of ODEs:
#
# $$ \frac{dx}{dt} = \sigma (y - x) $$
# $$ \frac{dy}{dt} = x (\rho - z) -y $$
# $$ \frac{dz}{dt} = xy - \beta z $$
#
# To define a system of differential equations in DifferentialEquations.jl, we define our `f` as a vector function with a vector initial condition. Thus, for the vector `u = [x,y,z]'`, we have the derivative function:

function lorenz!(du,u,p,t)
    σ,ρ,β = p
    du[1] = σ*(u[2]-u[1])
    du[2] = u[1]*(ρ-u[3]) - u[2]
    du[3] = u[1]*u[2] - β*u[3]
end
#----------------------------------------------------------------------------

# Notice here we used the in-place format which writes the output to the preallocated vector `du`. For systems of equations the in-place format is faster. We use the initial condition $u_0 = [1.0,0.0,0.0]$ as follows:

u0 = [1.0,0.0,0.0]
#----------------------------------------------------------------------------

# Lastly, for this model we made use of the parameters `p`. We need to set this value in the `ODEProblem` as well. For our model we want to solve using the parameters $\sigma = 10$, $\rho = 28$, and $\beta = 8/3$, and thus we build the parameter collection:

p = (10,28,8/3) ## we could also make this an array, or any other type!
#----------------------------------------------------------------------------

# Now we generate the `ODEProblem` type. In this case, since we have parameters, we add the parameter values to the end of the constructor call. Let's solve this on a time span of `t=0` to `t=100`:

tspan = (0.0,100.0)
prob = ODEProblem(lorenz!,u0,tspan,p)
#----------------------------------------------------------------------------

# Now, just as before, we solve the problem:

sol = solve(prob)
#----------------------------------------------------------------------------

# The same solution handling features apply to this case. Thus `sol.t` stores the time points and `sol.u` is an array storing the solution at the corresponding time points.
#
# However, there are a few extra features which are good to know when dealing with systems of equations. First of all, `sol` also acts like an array. `sol[i]` returns the solution at the `i`th time point.

sol.t[10],sol[10]
#----------------------------------------------------------------------------

# Additionally, the solution acts like a matrix where `sol[j,i]` is the value of the `j`th variable at time `i`:

sol[2,10]
#----------------------------------------------------------------------------

# We can get a real matrix by performing a conversion:

A = convert(Array,sol)
#----------------------------------------------------------------------------

# This is the same as sol, i.e. `sol[i,j] = A[i,j]`, but now it's a true matrix. Plotting will by default show the time series for each variable:

plot(sol)
#----------------------------------------------------------------------------

# If we instead want to plot values against each other, we can use the `vars` command. Let's plot variable `1` against variable `2` against variable `3`:

plot(sol,vars=(1,2,3))
#----------------------------------------------------------------------------

# This is the classic Lorenz attractor plot, where the `x` axis is `u[1]`, the `y` axis is `u[2]`, and the `z` axis is `u[3]`. Note that the plot recipe by default uses the interpolation, but we can turn this off:

plot(sol,vars=(1,2,3),denseplot=false)
#----------------------------------------------------------------------------

# Yikes! This shows how calculating the continuous solution has saved a lot of computational effort by computing only a sparse solution and filling in the values! Note that in vars, `0=time`, and thus we can plot the time series of a single component like:

plot(sol,vars=(0,2))
#----------------------------------------------------------------------------

# ### A DSL for Parameterized Functions
#
# In many cases you may be defining a lot of functions with parameters. There exists the domain-specific language (DSL) defined by the `@ode_def` macro for helping with this common problem. For example, we can define the Lotka-Volterra equation:
#
# $$ \frac{dx}{dt} = ax - bxy $$
# $$ \frac{dy}{dt} = -cy + dxy $$
#
# as follows:

function lotka_volterra!(du,u,p,t)
  du[1] = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = -p[3]*u[2] + p[4]*u[1]*u[2]
end
#----------------------------------------------------------------------------

# However, that can be hard to follow since there's a lot of "programming" getting in the way. Instead, you can use the `@ode_def` macro:

lv! = @ode_def LotkaVolterra begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d
#----------------------------------------------------------------------------

# We can then use the result just like an ODE function from before:

u0 = [1.0,1.0]
p = (1.5,1.0,3.0,1.0)
tspan = (0.0,10.0)
prob = ODEProblem(lv!,u0,tspan,p)
sol = solve(prob)
plot(sol)
#----------------------------------------------------------------------------

# Not only is the DSL convenient syntax, but it does some magic behind the scenes. For example, further parts of the tutorial will describe how solvers for stiff differential equations have to make use of the Jacobian in calculations. Here, the DSL uses symbolic differentiation to automatically derive that function:

lv!.Jex
#----------------------------------------------------------------------------

# The DSL can derive many other functions; this ability is used to speed up the solvers. An extension to DifferentialEquations.jl, [Latexify.jl](https://korsbo.github.io/Latexify.jl/latest/tutorials/parameterizedfunctions.html), allows you to extract these pieces as LaTeX expressions.

# ## Internal Types
#
# The last basic user-interface feature to explore is the choice of types. DifferentialEquations.jl respects your input types to determine the internal types that are used. Thus since in the previous cases, when we used `Float64` values for the initial condition, this meant that the internal values would be solved using `Float64`. We made sure that time was specified via `Float64` values, meaning that time steps would utilize 64-bit floats as well. But, by simply changing these types we can change what is used internally.
#
# As a quick example, let's say we want to solve an ODE defined by a matrix. To do this, we can simply use a matrix as input.

A  = [1. 0  0 -5
      4 -2  4 -3
     -4  0  0  1
      5 -2  2  3]
u0 = rand(4,2)
tspan = (0.0,1.0)
f(u,p,t) = A*u
prob = ODEProblem(f,u0,tspan)
sol = solve(prob)
#----------------------------------------------------------------------------

# There is no real difference from what we did before, but now in this case `u0` is a `4x2` matrix. Because of that, the solution at each time point is matrix:

sol[3]
#----------------------------------------------------------------------------

# In DifferentialEquations.jl, you can use any type that defines `+`, `-`, `*`, `/`, and has an appropriate `norm`. For example, if we want arbitrary precision floating point numbers, we can change the input to be a matrix of `BigFloat`:

big_u0 = big.(u0)
#----------------------------------------------------------------------------

# and we can solve the `ODEProblem` with arbitrary precision numbers by using that initial condition:

prob = ODEProblem(f,big_u0,tspan)
sol = solve(prob)
#----------------------------------------------------------------------------

sol[1,3]
#----------------------------------------------------------------------------

# To really make use of this, we would want to change `abstol` and `reltol` to be small! Notice that the type for "time" is different than the type for the dependent variables, and this can be used to optimize the algorithm via keeping multiple precisions. We can convert time to be arbitrary precision as well by defining our time span with `BigFloat` variables:
#

prob = ODEProblem(f,big_u0,big.(tspan))
sol = solve(prob)
#----------------------------------------------------------------------------

# Let's end by showing a more complicated use of types. For small arrays, it's usually faster to do operations on static arrays via the package [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl). The syntax is similar to that of normal arrays, but for these special arrays we utilize the `@SMatrix` macro to indicate we want to create a static array.

using StaticArrays
A  = @SMatrix [ 1.0  0.0 0.0 -5.0
                4.0 -2.0 4.0 -3.0
               -4.0  0.0 0.0  1.0
                5.0 -2.0 2.0  3.0]
u0 = @SMatrix rand(4,2)
tspan = (0.0,1.0)
f(u,p,t) = A*u
prob = ODEProblem(f,u0,tspan)
sol = solve(prob)
#----------------------------------------------------------------------------

sol[3]
#----------------------------------------------------------------------------

# ## Conclusion
#
# These are the basic controls in DifferentialEquations.jl. All equations are defined via a problem type, and the `solve` command is used with an algorithm choice (or the default) to get a solution. Every solution acts the same, like an array `sol[i]` with `sol.t[i]`, and also like a continuous function `sol(t)` with a nice plot command `plot(sol)`. The Common Solver Options can be used to control the solver for any equation type. Lastly, the types used in the numerical solving are determined by the input types, and this can be used to solve with arbitrary precision and add additional optimizations (this can be used to solve via GPUs for example!). While this was shown on ODEs, these techniques generalize to other types of equations as well.


using Distributions
using Plots
using QuadGK

Σ₁₁, Σ₁₂, Σ₁₃ = 1.0, 0.5, 0.1
Σ₂₁, Σ₂₂, Σ₂₃ = 0.5, 2.0, 0.2
Σ₃₁, Σ₃₂, Σ₃₃ = 0.1, 0.2, 0.5
μ₁=0
μ₂=-1.0
μ₃=2.0

μ = [
    μ₁
    μ₂
    μ₃
]

Σ = [
    Σ₁₁ Σ₁₂ Σ₁₃
    Σ₂₁ Σ₂₂ Σ₂₃
    Σ₃₁ Σ₃₂ Σ₃₃
]

μ_marginal = [
    μ₁
    μ₂
]

Σ_marginal = [
    Σ₁₁ Σ₁₂
    Σ₂₁ Σ₂₂
]

d = MvNormal(μ, Σ)
data = rand(d,250)

"""
p(x,y) = \\int p(x,y,z) dz
"""
function marginal(d, x, y)
    integral, err = quadgk(
        z->pdf(d,[x,y,z]),
        μ[2] - maximum(abs.(Σ)),
        μ[2] + maximum(abs.(Σ)),
        rtol = 1e-8,
    )
    return integral
end

marginal(y,z)=marginal(d,y,z)
contour(collect(-5:0.1:5), collect(-5:0.1:5), marginal)
scatter!(data[1,:], data[2,:], label="numeric")

# the pdf function p(x,y) should be 𝒩(μ_marginal, Σ_marginal).
d = MvNormal(μ_marginal, Σ_marginal)
data = rand(d,250)
scatter!(data[1,:],data[2,:],label="theoretial")


## Simple random graph generators ##

"""
    rand_geom(p::Real)

Sample a value from geometric distribution with probability of success `p`.
This code is adapted from the [Distributions.jl](https://github.com/JuliaStats/Distributions.jl)
package.
"""
rand_geom(p::Real) = floor(Int, -randexp() / log1p(-p))

"""
    gilbert_fill!(g::Graph, p::Real)

Fill graph by choosing each edge randomly with probability `p`.
"""
function gilbert_fill!(g::UndirectedGraph, p::Real)
    i, j = 1, 1
    n = nodecount(g)
    while true
        j += 1 + rand_geom(p)
        while j > n
            i += 1
            i > n && return
            j = j - n + i
        end
        addedge!(g, i, j)
    end
end

function gilbert_fill!(g::DirectedGraph, p::Real)
    i, j = 1, 1
    n = nodecount(g)
    while true
        j_ = j + 1 + rand_geom(p)
        j = (j < i && j_ >= i) ? j_ + 1 : j_
        while j > n
            i += 1
            i > n && return
            j -= n
            j += j >= i
        end
        addedge!(g, i, j)
    end
end

"""
    graph_gilbert(n::Integer, p::Real)

Generate a random graph according to Gilbert model with `n` nodes
and expected density `p`.

An independent Bernoulli trial is simulated for each possible edge in the graph
and it is added with probability `p`.
Note that in this case the number of edges is different each time, but the
matiematical expectation for the graph's density is `p`.

# References

Gilbert, E.N. (1959). "Random Graphs". Annals of Mathematical Statistics. 30 (4): 1141–1144.
doi:10.1214/aoms/1177706098.

Batagelj, V. and Brandes, U. (2005) "Efficient generation of large random networks".
Physical Review E. 71 (3).
"""
function graph_gilbert(n::Integer, p::Real; kwargs...)
    @assert(0 <= p <= 1, "edge probability (p) should be in [0, 1] interval")
    g = Graph(n; kwargs...)
    p > eps() && gilbert_fill!(g, p)
    return g
end

"""
    graph_erdos_renyi(n::Integer, m::Integer)

Generate an [Erdős–Rényi random graph](https://en.wikipedia.org/wiki/Erdős–Rényi_model)
with `n` nodes and `m` edges.

From mathematical viewpoint, the graph is picked at random with uniform probability
from a set of all possible graphs having `n` nodes and `m` edges.
In practice, an empty graph with `n` nodes is created first and
`m` edges get added between the randomly chosen pairs of its nodes.

# References

Erdős, P.; Rényi, A. (1959). "On Random Graphs. I". Publicationes Mathematicae. 6: 290–297.
"""
function graph_erdos_renyi(n::Integer, m::Integer; kwargs...)
    g = Graph(n; kwargs...)
    if m < maxedgecount(g)
        while edgecount(g) < m
            x, y = rand(1:n), rand(1:n)
            if x == y || hasedge(g, x, y)
                continue
            else
                addedge!(g, x, y)
            end
        end
    elseif m == maxedgecount(g)
        gilbert_fill!(g, 1.)
    else
        error("too many edges (m) for Erdős–Rényi graph with n nodes")
    end
    return g
end

@deprecate(erdosrenyi, graph_erdos_renyi)
@deprecate(graph_erdosrenyi, graph_erdos_renyi)

"""
    graph_random(n::Integer, m::Integer)
    graph_random(n::Integer, p::AbstractFloat)

A shorthand for creating random networks. There are two similar but distinct models:

* `(n, m)` model, also known as Erdős–Rényi model. It randomly selects a
  graph from a set of all possible graphs which have `n` nodes and `m` edges.

* `(n, p)` model, introduced by Gilbert. For a graph on `n` nodes,
  it includes each edge with probability `p` independently of the others.
  This gives a graph with mathematical expectation of density equal to `p`.

This method is just a shorthand for [`graph_erdos_renyi`](@ref) and
[`graph_gilbert`](@ref), which implement these models and are chosen automatically
based on the argument types.

If you need a graph of high density, use `(n, p)` model as it is much faster.
If you want an exact number of edges with lower density, use `(n, m)` model.

# Examples

```julia-repl
julia> g = graph_random(4, 1.)
4-node 12-edge directed multigraph

julia> g = graph_random(4, 1)
4-node 1-edge directed multigraph
```

# References

[Wikipedia](https://en.wikipedia.org/wiki/Erdős–Rényi_model)
"""
graph_random(n::Integer, m::Integer; kwargs...) =
    graph_erdos_renyi(n::Integer, m::Integer; kwargs...)
graph_random(n::Integer, p::AbstractFloat; kwargs...) =
    graph_gilbert(n::Integer, p::AbstractFloat; kwargs...)


include("header.jl")

#=
# Submerged surfaces
In this notebook, we will learn how to calculate force and moment on
flat surfaces immersed in static fluids
=#

# ### Set up the module
using MAE103
#-
using Plots

#=
# Example
A circular gate of 3 m diameter is located in an inclined wall (60 degrees
from horizontal) of a large reservoir filled with water. A horizontal shaft runs through
the center of the gate. The shaft is at a depth of 12 m. The other side
of the gate is exposed to air at atmospheric pressure.
* What is the force on the gate?
* Where does this force effectively act?
* What is the moment applied by the water on the gate shaft?
=#
#-
#=
Here, we are working with a simple, symmetric shape (a circle), and we
know that the centroid of a circle is at its center. Since the shaft runs
through the center, we know that the centroid is at a depth of 12 m.
=#
hc = Depth(12)
d = Diameter(3)
Θ = 60*π/180

#=
The force is easy to compute, since it is just the specific weight times
area times depth.

$$F_R = \gamma h_c A$$

(The atmospheric pressure contributes nothing in this problem, since it
acts the same on both sides.)
=#
A = Area(Circle,d/2)
FR = Force(SpecificWeight(Water)*A*hc)

#=
To find the center of pressure, we use the formula

$$y_R = \dfrac{I_{xc}}{y_c A} + y_c$$

This the location along the $y$ axis, which runs from the free surface downward, parallel
to the plate. In the formula, $I_{xc}$ is the second area moment about the $x$ axis running
through the centroid of the plate, and $y_c$ is the $y$ coordinate of the centroid.

But $y_c$ is simply $h_c/\sin\theta$:
=#
yc = Length(hc/sin(Θ))

#=
and $I_{xc}$ we can compute with a handy function:
=#
Ixc = SecondAreaMomentX(Circle,d/2)

#=
Now put them together to get the center of pressure
=#
yR = Length(Ixc/(yc*A) + yc)

#=
How far is this from the centroid?
=#
yR - yc

#=
Remember, it is lower than the centroid, because there are higher pressures
acting below the centroid than above it.

To determine the moment $M_R$ about the shaft, we simply need to remember that the
force is acting at the center of pressure, so the moment is equal to the force times the distance
of this center of pressure from the shaft (i.e., from the centroid)
=#
MR = Moment((yR - yc)*FR)

#=
#### A follow-up question
What happens to these values as the plate gets deeper?

Obviously force gets larger. What about $y_R$ and $M_R$? Let's set the depth
at 20 m, and recalculate:
=#
hc = Depth(20)
FR = Force(SpecificWeight(Water)*A*hc)
yc = Length(hc/sin(Θ))
yR = Length(Ixc/(yc*A) + yc)
yR - yc
#=
so it gets a little closer to the centroid as the plate gets lower. What about
the moment?
=#
MR = Moment((yR - yc)*FR)
#=
It's exactly the same! In fact, **the moment about the centroid does not
depend on how deep the plate is**.
=#

#=
#### Another follow-up question
What happens if the plate becomes more vertical? We'll set $h_c$ to 12 m again
and make $Θ = 90^\circ$. Only $y_c$ changes:
=#
hc = Depth(12)
Θ = 90*π/180
FR = Force(SpecificWeight(Water)*A*hc)
#=
Note that **force is not affected** by a change in angle, since only depth
of the centroid matters.
=#
yc = Length(hc/sin(Θ))
yR = Length(Ixc/(yc*A) + yc)
yR - yc

#=
So the center of pressure is a little further now from the centroid than
when it was at 60 degrees.

The moment:
=#
MR = Moment((yR - yc)*FR)

#=
So **moment gets larger as the surface becomes more vertical**. This is because the
center of pressure has moved away from the centroid. This is intuitive if you
think about the other extreme: a completely horizontal surface. In that situation,
the pressure acts uniformly across the plate, so there is **zero net moment**
on a horizontal surface.
=#


using NewtonMethod2
using Test

@testset "NewtonMethod2.jl" begin

### 2 & 3 & 4
f(x)=x^2+2*x-3
f_prime(x)=2*x+2
x0=2.0
t2=newtonroot(f, f_prime; x0, tol=1E-7, maxiter=1000)[1]
t3=newtonroot2(f; x0, tol=1E-7, maxiter=1000)[1]
t4=newtonroot(f, f_prime; x0=BigFloat(x0), tol=1E-7, maxiter=1000)[1]
t5=newtonroot2(f; x0=BigFloat(x0), tol=1E-7, maxiter=1000)[1]
@test t2 ≈ 1.0
@test t3 ≈ 1.0
@test t4 ≈ 1.0
@test t5 ≈ 1.0

f(x)=x^2+2*x-8
f_prime(x)=2*x+2
x0=1.0
t2=newtonroot(f, f_prime; x0, tol=1E-7, maxiter=1000)[1]
t3=newtonroot2(f; x0, tol=1E-7, maxiter=1000)[1]
t4=newtonroot(f, f_prime; x0=BigFloat(x0), tol=1E-7, maxiter=1000)[1]
t5=newtonroot2(f; x0=BigFloat(x0), tol=1E-7, maxiter=1000)[1]
@test t2 ≈ 2.0
@test t3 ≈ 2.0
@test t4 ≈ 2.0
@test t5 ≈ 2.0

f(x)=x-2
f_prime(x)=1
x0=0.0
t2=newtonroot(f, f_prime; x0, tol=1E-7, maxiter=1000)[1]
t3=newtonroot2(f; x0, tol=1E-7, maxiter=1000)[1]
t4=newtonroot(f, f_prime; x0=BigFloat(x0), tol=1E-7, maxiter=1000)[1]
t5=newtonroot2(f; x0=BigFloat(x0), tol=1E-7, maxiter=1000)[1]
@test t2 ≈ 2.0
@test t3 ≈ 2.0
@test t4 ≈ 2.0
@test t5 ≈ 2.0

### 5
f(x)=2+x^2
f_prime(x)=2*x
x0=1.0
t2=newtonroot(f, f_prime; x0, tol=1E-7, maxiter=1000)[1]
t3=newtonroot2(f; x0, tol=1E-7, maxiter=1000)[1]
@test t2 == nothing
@test t3 == nothing

### 6
f(x)=x^2+2*x-3
f_prime(x)=2*x+2
x0=2.0
t2=newtonroot(f, f_prime; x0, tol=1E-7, maxiter=3)[1]
t3=newtonroot2(f; x0, tol=1E-7, maxiter=3)[1]
@test t2 == nothing
@test t3 == nothing

### 7
f(x)=x^2+2*x-3
f_prime(x)=2*x+2
x0=2.0
t2=newtonroot(f, f_prime; x0, tol=1E-7, maxiter=1000)[2]
t3=newtonroot2(f; x0, tol=1E-7, maxiter=1000)[2]
@test t2 <= 1E-7
@test t3 <= 1E-7 ###

end


module TestSphere

include("../src/Sphere.jl")
import .Sphere

using Test

const tol = 5e-13

export test_swal_inner_product, test_norm_swal_lap, test_norm_filter

"""
Computes inner product between different Y^s_{lm}

test_swal_inner_product(
   ny::Int64,
   spin::Int64,
   m_ang::Int64
   )::Nothing
"""
function test_swal_inner_product(
      ny::Int64,
      spin::Int64,
      m_ang::Int64
   )::Nothing

   nl = Sphere.num_l(ny)

   vals = Sphere.swal_vals(ny,spin,m_ang)

   println("Test: Y^s_{lm} orthogonality:\tny=$ny\tspin=$spin\tm_ang=$m_ang\tl=[0..$nl]")
   for i=1:nl
      for j=i:nl
         val = Sphere.inner_product(vals[:,i],vals[:,j])
         if i==j
            @test abs(val-1.0) < tol
         else
            @test abs(val) < tol
         end
      end
   end
   return nothing
end

"""
Computes the norm of the difference between the numerical and
exact spin-weighted spherical laplacian for a given
spin-weighted spherical harmonic.

test_norm_swal_lap(
   ny::Int64,
   spin::Int64,
   m_ang::Int64,
   l_ang::Int64
   )::Nothing
"""
function test_norm_swal_lap(
      ny::Int64,
      spin::Int64,
      m_ang::Int64,
      l_ang::Int64
   )::Nothing

   Yv = Sphere.Y_vals(ny)

   swal = [Sphere.swal(spin,m_ang,l_ang,y) for y in Yv]

   swal_lap_v1 = [-(l_ang-spin)*(l_ang+spin+1.0)*v for v in swal]
   swal_lap_v2 = zeros(Float64,ny)

   lap = Sphere.swal_laplacian_matrix(ny,spin,m_ang)
   for j=1:ny
      for k=1:ny
         swal_lap_v2[j] += swal[k]*lap[k,j]
      end
   end

   ## compute integral over interval [-1,1] of difference
   n = Sphere.inner_product(ones(ny),swal_lap_v1 .- swal_lap_v2)

   println("Test: Δ_{S2}Y^s_{lm}:\tny=$ny\tspin=$spin\tm_ang=$m_ang\tl=$l_ang")
   @test abs(n) < tol

   return nothing
end

"""
Test that the swal filter matrix acts as a low pass filter
for angular perturbations.
"""
function test_norm_filter(
      ny::Int64,
      spin::Int64,
      m_ang::Int64,
      l_ang::Int64
   )::Nothing

   Yv = Sphere.Y_vals(ny)

   ## essentially random points multiplied by spherical harmonic;
   ## with the filter the norm should go to zero
   v1 = [0.1*(rand()-0.5)*Sphere.swal(spin,m_ang,l_ang,y) for y in Yv]
   v2 = zeros(Float64,ny)

   filter = Sphere.swal_filter_matrix(ny,spin,m_ang)

   for j=1:ny
      for k=1:ny
         v2[j] += v1[k]*filter[k,j]
      end
   end

   ## compute integral over interval [-1,1] of difference
   i1 = abs(Sphere.inner_product(ones(ny), abs.(v1)))
   i2 = abs(Sphere.inner_product(ones(ny), abs.(v2)))

   println("Test: filter(rand) smoothens:\tny=$ny\tspin=$spin\tm_ang=$m_ang\tl_ang=$l_ang")
   @test i2 < i1

   return nothing
end

end


import Distributions: quantile, Normal

"""
```julia
p₁ = binom_confidence(n::Int,c::Real)
```
Returns the probability `p₁` of successes in the binomial distribution (p=1/2) of
`n` trials with confidence `c`.

# Example
At c=0.95, i.e. 95% confidence, n=1000 tosses of
a coin will yield not more than
```julia
julia> p₁ = BitInformation.binom_confidence(1000,0.95)
0.5309897516152281
```
about 53.1% heads (or tails)."""
function binom_confidence(n::Int,c::Real)
    return 0.5 + quantile(Normal(),1-(1-c)/2)/(2*sqrt(n))
end

"""
```julia
Hf = binom_free_entropy(n::Int,c::Real,base::Real=2)
```
Returns the free entropy `Hf` associated with `binom_confidence`."""
function binom_free_entropy(n::Int,c::Real,base::Real=2)
    p = binom_confidence(n,c)
    return 1 - entropy([p,1-p],base)
end


function fixpnt(δ=1e-3, x0=1.0)
    x = x0
    iter = 0
    while true
        iter += 1
        x0 = x
        x = atan(2x0)
        if abs(x-x0) ≤ δ
            println("Terminating after $iter iterations, x = $x")
            return
        end
    end
end

fixpnt()

fixpnt(1e-10)

fixpnt(1e-10, 0.1)

function fixpnt2(δ=1e-3, x0=1.0; maxiter=10)
    x = x0
    iter = 0
    while true
        iter += 1
        x0 = x
        x = atan(2x0)
        if iter ≥ maxiter || abs(x-x0) ≤ δ
            println("Terminating after $iter iterations, x = $x")
            return
        end
    end
end

fixpnt2(1e-10, maxiter=15)


export fibonacci_bound
export fibonacci_count
export unique_products_of_ndigits
export multiset_combination!
export permutation_matrix!

"""
Returns an array of the Fibonacci numbers up to a bound n.
"""
function fibonacci_bound(n::BigInt=1000)::Array{BigInt}
    a = BigInt(1)
    b = BigInt(2)
    fibs = zeros(BigInt, 1)
    fibs[1] = 1
    while a < n
        (a, b) = (b, a + b)
        push!(fibs, a)
    end
    return fibs
end

"""
Returns an array of the first n Fibonacci numbers.
"""
function fibonacci_count(n::Integer=20)::Array{BigInt}
    a = BigInt(1)
    b = BigInt(2)
    fibs = zeros(BigInt, n)
    fibs[1] = 1
    for i in 2:n
        (a, b) = (b, a + b)
        @inbounds fibs[i] = a
    end
    return fibs
end

"""
Return a list of products of numbers with "ndigit" digits.
"""
function unique_products_of_ndigits(n_digits::Integer)::Array{Integer}
    maxnum = 10^n_digits - 1
    minnum = 10^(n_digits - 1)
    return [i * j for i in minnum:maxnum for j in minnum:maxnum if i < j]
end

"""
Return a multiset (combinations with replacement, order doesn't matter).  Uses
a recursive solution to build up a full set of indexes.
"""
function multiset_combination!(storage::Array{Int8,2},
                               digit_index::Integer,
                               index::Integer,
                               rstart::Integer,
                               rend::Integer)
    if digit_index == 0
        storage[:, index + 1] = storage[:, index]
        return index + 1
    end

    for i in rstart:rend
        storage[digit_index, index] = i
        index = multiset_combination!(storage, digit_index - 1, index, i, rend)
    end
    return index
end

"""
Create array of permutations for which valid_func is true.
"""
function permutation_matrix!(elementlist::Array, matrix::Array, valid_func, k::Integer=1)
    if k == length(elementlist)
        if valid_func(elementlist)
            append!(matrix, elementlist)
        end
    else
        for i in k:length(elementlist)
            elementlist[k], elementlist[i] = elementlist[i], elementlist[k]
            permutation_matrix!(elementlist, matrix, valid_func, k + 1)
            elementlist[k], elementlist[i] = elementlist[i], elementlist[k]
        end
    end
end


using NanStats
using Printf

@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

vec1=[NaN,3,4,NaN]
mn=nanmean(vec1);
varn=nanvar(vec1);
kurn=nankurt(vec1);
sken=nanskew(vec1);

print("vec1=$(vec1)\n");

@testset "nanmean" begin
	@test mn == 3.5
        print("The mean for vec1 ignoring NaN values is $(mn)\n");
end

@testset "nanvar" begin
	@test varn == 0.25
        print("The variance for vec1 ignoring NaN values is $(varn)\n")
end

@testset "nankurt" begin
	@test kurn == 1.0
        print("The kurtosis for vec1 ignoring NaN values is $(kurn)\n")
end

@testset "nanskew" begin
	@test sken == 0.0
        print("The skewness for vec1 ignoring NaN values is $(sken)\n")
end


using SpecialFunctions: gamma

"
Evaluate Jacobi polynomial of type
(alpha, beta) > -1, (alpha+beta != -1)
at points x for order N. Returns array of values, same length as x

See: https://en.wikipedia.org/wiki/Jacobi_polynomials

NOTE: This calculates the normalised Jacobi polynomial

"
function JacobiP(x, alpha, beta, N::Int)

    nvalues = length(x)

    PL = zeros(Float64, nvalues, N+1)

    gamma0 = 2^(alpha+beta+1)/(alpha+beta+1) * gamma(alpha+1) * gamma(beta+1)/gamma(alpha+beta+1)
    PL[:,1] .= 1.0 / sqrt(gamma0)

    if N == 0
        return PL[:,1]
    end

    gamma1 = (alpha+1)*(beta+1)/(alpha+beta+3)*gamma0
    PL[:,2] = ((alpha+beta+2)*x/2 .+ (alpha-beta)/2)/sqrt(gamma1)

    if N == 1
        return PL[:,2]
    end

    aold = 2/(2+alpha+beta)*sqrt((alpha+1)*(beta+1)/(alpha+beta+3))

    # Forward recurrence
    for i = 1:(N-1)
        h1 = 2*i + alpha + beta
        anew = 2/(h1+2)*sqrt( (i+1)*(i+1+alpha+beta)*(i+1+alpha)*(i+1+beta)/(h1+1)/(h1+3))
        bnew = - (alpha^2 - beta^2)/h1/(h1+2)
        PL[:,i+2] = (1/anew)*(-aold * PL[:,i] + (x .- bnew) .* PL[:,i+1])
        aold = anew
    end
    return PL[:,end]
end

"Evaluate the derivative of the Jacobi polynomial
of type (alpha,beta) > -1 at points r for order N

## Returns
array of the same length as r
"
function GradJacobiP(r, alpha, beta, N::Int)
    if N == 0
        zeros(length(r))
    else
        sqrt(N * (N+alpha+beta+1))*JacobiP(r, alpha+1, beta+1, N-1)
    end
end

struct JacobiBasis1D
    order::Int  # Number of polynomials
end

function basisFunction(basis::JacobiBasis1D, locations, index::Int)
    JacobiP(locations, 0, 0, index-1)
end

function basisGradient(basis::JacobiBasis1D, locations, index::Int)
    GradJacobiP(locations,0,0,index-1)
end


# p4.jl - periodic spectral differentiation

# Set up grid and differentiation matrix:
N = 24; h = 2*pi/N; x = h*(1:N);
column = [0; @. .5*(-1)^(1:N-1)*cot((1:N-1)*h/2)];
D = toeplitz(column,column[[1;N:-1:2]]);

# Differentiation of a hat function:
v = @. max(0,1-abs(x-pi)/2); clf();
subplot(2,2,1);  plot(x,v,".-",markersize=6);
axis([0,2pi,-.5,1.5]); grid(true); title("function");
subplot(2,2,2), plot(x,D*v,".-",markersize=6);
axis([0,2pi,-1,1]); grid(true); title("spectral derivative");

# Differentiation of exp(sin(x)):
v = @. exp(sin(x)); vprime = @. cos(x)*v;
subplot(2,2,3), plot(x,v,".-",markersize=6);
axis([0,2pi,0,3]), grid(true);
subplot(2,2,4), plot(x,D*v,".-",markersize=6);
axis([0,2pi,-2,2]), grid(true);
error = signif(norm(D*v-vprime,Inf),5);
text(2.2,1.4,"max error = $error",fontsize=8);


# ---
# title: 258. Add Digits
# id: problem258
# author: zhwang
# date: 2022-03-03
# difficulty: Easy
# categories: Math
# link: <https://leetcode.com/problems/add-digits/description/>
# hidden: true
# ---
#
# Given a non-negative integer `num`, repeatedly add all its digits until the
# result has only one digit.
#
# **Example:**
#
#
#
#     Input: 38
#     Output: 2
#     Explanation: The process is like: 3 + 8 = 11, 1 + 1 = 2.
#                  Since 2 has only one digit, return it.
#
#
# **Follow up:**
# Could you do it without any loop/recursion in O(1) runtime?
#
#
## @lc code=start
using LeetCode

function add_digits(num::Int)::Int
    num <= 9 && return num
    return iszero(num % 9) ? 9 : num % 9
end
## @lc code=end

## @lc test=start
@testset "258.add-digits.jl" begin
    @test all(add_digits(i) == i for i in 0:9)
    @test add_digits(18) == 9
    @test add_digits(10) == 1
    @test add_digits(38) == 2
end
## @lc test=end

## @lc info=start
# link: [solution to 258]()
## @lc info=end


#= Fornberg algorithm

Implements the Fornberg (1988) algorithm
(https://doi.org/10.1090/S0025-5718-1988-0935077-0), adapted from the
DiffEqOperators module.
=#
"""
    calculate_weights(m, n, s)

Return finite difference weights. Notation follows that of Fornberg, used in
Mathematica as well:
https://reference.wolfram.com/language/tutorial/NDSolveMethodOfLines.html

# Arguments
* `m::Int`: order of the derivative
* `n::Int`: number of grid intervals enclosed in the stencil
* `s::Int`: number of grid intervals between the point at which the derivative
            is approximated and the leftmost edge of the stencil. centred
            formulas always have s=n/2.
"""
function calculate_weights(m::Int, n::Int, s::Int)
    stencil_length = n + 1

    # dummy array with the relative coordinates of the stencil (eg. central
    # differences need coordinates centred at 0)
    x  = -s : (stencil_length - s - 1)
    # point in the array 'x' for which we need the coefficients
    x0 = x[1] + s

    N = length(x)
    @assert m < N "Not enough points for the requested order."

    c1 = 1
    c4 = x[1] - x0
    C = zeros(Rational{Int}, N, m+1)
    C[1,1] = 1
    @inbounds for i in 1:N-1
        i1 = i + 1
        mn = Base.min(i, m)
        c2 = 1
        c5 = c4
        c4 = x[i1] - x0
        for j in 0:i-1
            j1 = j + 1
            c3 = x[i1] - x[j1]
            c2 *= c3
            if j == i-1
                for s in mn:-1:1
                    s1 = s + 1
                    C[i1,s1] = c1*(s*C[i,s] - c5*C[i,s1]) // c2
                end
                C[i1,1] = -c1*c5*C[i,1] // c2
           end
            for s in mn:-1:1
                s1 = s + 1
                C[j1,s1] = (c4*C[j1,s1] - s*C[j1,s]) // c3
            end
            C[j1,1] = c4 * C[j1,1] // c3
        end
        c1 = c2
    end

    #=
    Fix the problem of numerical instability which occurs when the sum of the
    stencil_coefficients is not exactly 0:
https://scicomp.stackexchange.com/questions/11249/numerical-derivative-and-finite-difference-coefficients-any-update-of-the-fornb

    Stack Overflow answer on this issue:
    http://epubs.siam.org/doi/pdf/10.1137/S0036144596322507 - Modified Fornberg Algorithm
    =#

    _C = C[:,end]
    _C[div(N,2)+1] -= sum(_C)
    return _C
end


# Wilkinson-type polynomial defined by its roots:
# wₙ(x) = (x-1)⋅(x-2)⋅ ⋯ ⋅ (x-n)

using Polynomials

# function generate_wilkinson(n)#, T=BigFloat)   # SLOW

#     p = poly(collect(1:n))
#     p_prime = polyder(p)

#     coeffs = p.a #[convert(T, coeff) for coeff in p.a]
#     # coeffs_prime = [convert(T, coeff) for coeff in p_prime.a]

#     function f(x)
#         total = coeffs[1]
#         for i in 2:length(coeffs)
#             total += coeffs[i]*x^(i-1)
#         end
#         total
#     end

#     f
# end

function generate_wilkinson_horner(n)
    p = poly(collect(1:n))  # make polynomial with roots 1 to n

    coeffs = p.a

    expr = :($(coeffs[end]))

    for i in length(coeffs)-1:-1:1
        expr = :($expr*x + $(coeffs[i]))
    end

    expr
end

function subscriptify(n::Int)
    subscript_digits = [c for c in "₀₁₂₃₄₅₆₇₈₉"]
    dig = reverse(digits(n))
    join([subscript_digits[i+1] for i in dig])
end

# Generate Wilkinson functions W₃ etc.:
for n in 1:15
    fn_name = symbol(string("W", subscriptify(n)))
    expr = generate_wilkinson_horner(n)

    @eval $(fn_name)(x) = $(expr)
end


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

A row of five black square tiles is to have a number of its tiles replaced with
coloured oblong tiles chosen from red (length two), green (length three), or
blue (length four).

If red tiles are chosen there are exactly seven ways this can be done. If green
tiles are chosen there are three ways. And if blue tiles are chosen there are
two ways.Assuming that colours cannot be mixed there are 7 + 3 + 2 = 12 ways of
replacing the black tiles in a row measuring five units in length.

How many different ways can the black tiles in a row measuring fifty units in
length be replaced if colours cannot be mixed and at least one coloured tile
must be used

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

using Memoize
using Base.Test

@memoize function ncases_swb(blocklen, rowlen)
    # Return the number of cases _s_tarting _w_ith a _b_lock

    if rowlen < blocklen
        return 0
    end

    result = 1
    # case for which there is only one block

    rowlen -= blocklen
    # remove room taken by the first block

    for nspace = 0:rowlen-blocklen
        # nbr spaces allowed between first block and second one

        result += ncases_swb(blocklen, rowlen-nspace)
    end

    result
end

function ncases(blocklen, rowlen)
    result = 0
    for nspace = 0:rowlen-blocklen
        # nbr spaces before first block
        result += ncases_swb(blocklen, rowlen-nspace)
    end

    result
end

function euler(rowlen)
    ncases(2, rowlen) + ncases(3, rowlen) + ncases(4, rowlen)
end

@test ncases(2, 5) == 7
@test ncases(3, 5) == 3
@test ncases(4, 5) == 2
@test euler(5) == 12

@printf("Euler 0116: %d", euler(50))


"""
    lsnormal(A,b)

Solve a linear least squares problem by the normal equations.
Returns the minimizer of ||b-Ax||.
"""
function lsnormal(A,b)
    N = A'*A;  z = A'*b;
    R = cholesky(N).U
    w = forwardsub(R',z)                   # solve R'z=c
    x = backsub(R,w)                       # solve Rx=z
    return x
end

"""
    lsqrfact(A,b)

Solve a linear least squares problem by QR factorization. Returns
the minimizer of ||b-Ax||.
"""
function lsqrfact(A,b)
    Q,R = qr(A)
    c = Q'*b
    x = backsub(R,c)
    return x
end

"""
    qrfact(A)

QR factorization by Householder reflections. Returns Q and R.
"""
function qrfact(A)
    m,n = size(A)
    Qt = diagm(ones(m))
    R = float(copy(A))
    for k in 1:n
        z = R[k:m,k]
        w = [ -sign(z[1])*norm(z) - z[1]; -z[2:end] ]
        nrmw = norm(w)
        if nrmw < eps() continue; end    # skip this iteration
        v = w / nrmw;
        # Apply the reflection to each relevant column of R and Q
        for j in k:n
            R[k:m,j] -= v*( 2*(v'*R[k:m,j]) )
        end
        for j in 1:m
            Qt[k:m,j] -= v*( 2*(v'*Qt[k:m,j]) )
        end
    end
    return Qt',triu(R)
end


"""
    non_relativistic_orbital(n, ℓ[, Z = 1])

Returns a callable object to evaluate the _reduced_ radial orbital ``P_{nℓ}(r)`` of a
hydrogen-like atom in non-relativistic theory, given by

```math
P_{nℓ}(r) =
\\sqrt{
    \\left(\\frac{2Z}{n}\\right)^3
    \\frac{(n-\\ell-1)!}{2n(n+\\ell)!}
}
\\left(\\frac{2Z}{n} \\cdot r\\right)^\\ell
L^{(2\\ell+1)}_{n-\\ell-1}\\left(\\frac{2Z}{n} \\cdot r\\right)
\\exp\\left(-\\frac{Z}{n} \\cdot r\\right)
```

``P_{nℓ}(r)`` is defined through

```math
\\psi_{nlm}(r,\\theta,\\varphi) = \\frac{1}{r} P_{nℓ}(r) Y_{ℓm}(\\theta,\\varphi)
```

where ``\\psi_{nlm}`` are the normalized solutions to the non-relativistic single-particle
Schrödinger equation for a point-like nucleus with charge `Z`, and where ``Y_{ℓm}`` are the
spherical harmonics.
"""
function non_relativistic_orbital(n::Integer, ℓ::Integer, Z = 1)
    n > 0 || throw(ArgumentError("Invalid n = $n, must be n ≥ 1"))
    0 ≤ ℓ < n || throw(ArgumentError("Invalid ℓ = $ℓ for n = $n, must be 0 ≤ ℓ < n"))
    # Normalization factor.
    # Note: n! = Γ(n+1), so (n-ℓ-1)!/(2n(n+ℓ)!) = Γ(n+1)/(2nΓ(n+ℓ+1))
    N = sqrt((2Z/n)^3*gamma(n-ℓ)/(2n*gamma(n+ℓ+1)))
    # Generalized Laguerre polynomial L^{(2ℓ+1)}_{n-ℓ-1}(q), where q = 2Zr/n:
    L = convert(Polynomial, Laguerre{2ℓ+1}([zeros(n-ℓ-1); 1], :q))
    # Single Polynomial object for the polynomial part, multiplied by r to get the
    # reduced polynomial:
    C = let r = Polynomial([0, 1], :r)
        N * (2*Z*r/n)^ℓ * L(2*Z*r/n) * r
    end
    # Multiply with the exponential exp(Zr/n) and return a callable object:
    q = -Z/n # pre-compute the coefficient in the exponential
    r -> C(r) * exp(q*r)
end

export non_relativistic_orbital


module ModuleB

using LinearAlgebra

export myPower, myDeflation, myPowerMethod, myTridiag, myTridiagX, myTridiagG, myTridEigQR,
mySymEigQR, stedc!, myJacobi, myLanczos, myBidiagX, myBidiagY, myBidiagG, myJacobiR

function myPower(A::Matrix,x::Vector,tol::Real)
    y=A*x
    ν=x⋅y
    steps=1
    while norm(y-ν*x)>tol
        x=y/norm(y)
        y=A*x
        ν=x⋅y
        steps+=1
    end
    ν, y/norm(y), steps
end

function myDeflation(A::Matrix,x::Vector)
    X,R=qr(x[:,:])
    # To make sure the returned matrix symmetric use
    # full(Symmetric(X[:,2:end]'*A*X[:,2:end]))
    X[:,2:end]'*A*X[:,2:end]
end

function myPowerMethod(A::Matrix{T}, tol::Real) where T
    n=size(A,1)
    S =  T==Int ? Float64 : T
    λ=Vector{S}(undef,n)
    for i=1:n
        λ[i],x,steps=myPower(A,ones(n-i+1),tol)
        A=myDeflation(A,x)
    end
    λ
end

function myTridiag(A::Matrix)
    # Normalized Householder vectors are stored in the lower
    # triangular part of A below the first subdiagonal
    n=size(A,1)
    T=Float64
    A=map(T,A)
    v=Vector{T}(undef,n)
    Trid=SymTridiagonal(zeros(n),zeros(n-1))
    for j = 1 : n-2
        μ = sign(A[j+1,j])*norm(A[j+1:n, j])
        if μ != zero(T)
            β =A[j+1,j]+μ
            v[j+2:n] = A[j+2:n,j] / β
        end
        A[j+1,j]=-μ
        A[j,j+1]=-μ
        v[j+1] = one(T)
        γ = -2 / (v[j+1:n]⋅v[j+1:n])
        w = γ* A[j+1:n, j+1:n]*v[j+1:n]
        q = w + γ * v[j+1:n]*(v[j+1:n]⋅w) / 2
        A[j+1:n, j+1:n] = A[j+1:n,j+1:n] + v[j+1:n]*q' + q*v[j+1:n]'
        A[j+2:n, j] = v[j+2:n]
    end
    SymTridiagonal(diag(A),diag(A,1)), tril(A,-2)
end

# Extract X
function myTridiagX(H::Matrix)
    n=size(H,1)
    T=Float64
    X = Matrix{T}(I,n,n)
    v=Vector{T}(undef,n)
    for j = n-2 : -1 : 1
        v[j+1] = one(T)
        v[j+2:n] = H[j+2:n, j]
        γ = -2 / (v[j+1:n]⋅v[j+1:n])
        w = γ * X[j+1:n, j+1:n]'*v[j+1:n]
        X[j+1:n, j+1:n] = X[j+1:n, j+1:n] + v[j+1:n]*w'
    end
    X
end

# Tridiagonalization using Givens rotations
function myTridiagG(A::Matrix)
    n=size(A,1)
    T=Float64
    X=Matrix{T}(I,n,n)
    for j = 1 : n-2
        for i = j+2 : n
            G,r=givens(A,j+1,i,j)
            A=(G*A)*G'
            X*=G'
        end
    end
    SymTridiagonal(diag(A),diag(A,1)), X
end

function myTridEigQR(A1::SymTridiagonal)
    A=deepcopy(A1)
    n=length(A.dv)
    T=Float64
    λ=Vector{T}(undef,n)
    B=Matrix{T}
    if n==1
        return map(T,A.dv)
    end
    if n==2
        τ=(A.dv[end-1]-A.dv[end])/2
        μ=A.dv[end]-A.ev[end]^2/(τ+sign(τ)*sqrt(τ^2+A.ev[end]^2))
        # Only rotation
        B=A[1:2,1:2]
        G,r=givens(B-μ*I,1,2,1)
        B=(G*B)*G'
        return diag(B)[1:2]
    end
    steps=1
    k=0
    while k==0 && steps<=10
        # Shift
        τ=(A.dv[end-1]-A.dv[end])/2
        μ=A.dv[end]-A.ev[end]^2/(τ+sign(τ)*sqrt(τ^2+A.ev[end]^2))
        # First rotation
        B=A[1:3,1:3]
        G,r=givens(B-μ*I,1,2,1)
        B=(G*B)*G'
        A.dv[1:2]=diag(B)[1:2]
        A.ev[1:2]=diag(B,-1)
        bulge=B[3,1]
        # Bulge chasing
        for i = 2 : n-2
            B=A[i-1:i+2,i-1:i+2]
            B[3,1]=bulge
            B[1,3]=bulge
            G,r=givens(B,2,3,1)
            B=(G*B)*G'
            A.dv[i:i+1]=diag(B)[2:3]
            A.ev[i-1:i+1]=diag(B,-1)
            bulge=B[4,2]
        end
        # Last rotation
        B=A[n-2:n,n-2:n]
        B[3,1]=bulge
        B[1,3]=bulge
        G,r=givens(B,2,3,1)
        B=(G*B)*G'
        A.dv[n-1:n]=diag(B)[2:3]
        A.ev[n-2:n-1]=diag(B,-1)
        steps+=1
        # Deflation criterion
        k=findfirst(abs.(A.ev) .< sqrt.(abs.(A.dv[1:n-1].*A.dv[2:n]))*eps(T))
        k=k==nothing ? 0 : k
    end
    λ[1:k]=myTridEigQR(SymTridiagonal(A.dv[1:k],A.ev[1:k-1]))
    λ[k+1:n]=myTridEigQR(SymTridiagonal(A.dv[k+1:n],A.ev[k+1:n-1]))
    λ
end

function myTridEigQR(A1::SymTridiagonal,U::Matrix)
    # U is either the identity matrix or the output from myTridiagX()
    A=deepcopy(A1)
    n=length(A.dv)
    T=Float64
    λ=Vector{T}(undef,n)
    B=Matrix{T}
    if n==1
        return map(T,A.dv), U
    end
    if n==2
        τ=(A.dv[end-1]-A.dv[end])/2
        μ=A.dv[end]-A.ev[end]^2/(τ+sign(τ)*sqrt(τ^2+A.ev[end]^2))
        # Only rotation
        B=A[1:2,1:2]
        G,r=givens(B-μ*I,1,2,1)
        B=(G*B)*G'
        U*=G'
        return diag(B)[1:2], U
    end
    steps=1
    k=0
    while k==0 && steps<=10
        # Shift
        τ=(A.dv[end-1]-A.dv[end])/2
        μ=A.dv[end]-A.ev[end]^2/(τ+sign(τ)*sqrt(τ^2+A.ev[end]^2))
        # First rotation
        B=A[1:3,1:3]
        G,r=givens(B-μ*I,1,2,1)
        B=(G*B)*G'
        U[:,1:3]*=G'
        A.dv[1:2]=diag(B)[1:2]
        A.ev[1:2]=diag(B,-1)
        bulge=B[3,1]
        # Bulge chasing
        for i = 2 : n-2
            B=A[i-1:i+2,i-1:i+2]
            B[3,1]=bulge
            B[1,3]=bulge
            G,r=givens(B,2,3,1)
            B=(G*B)*G'
            U[:,i-1:i+2]=U[:,i-1:i+2]*G'
            A.dv[i:i+1]=diag(B)[2:3]
            A.ev[i-1:i+1]=diag(B,-1)
            bulge=B[4,2]
        end
        # Last rotation
        B=A[n-2:n,n-2:n]
        B[3,1]=bulge
        B[1,3]=bulge
        G,r=givens(B,2,3,1)
        B=(G*B)*G'
        U[:,n-2:n]*=G'
        A.dv[n-1:n]=diag(B)[2:3]
        A.ev[n-2:n-1]=diag(B,-1)
        steps+=1
        # Deflation criterion
        k=findfirst(abs.(A.ev) .< sqrt.(abs.(A.dv[1:n-1].*A.dv[2:n]))*eps())
        k=k==nothing ? 0 : k
    end
    λ[1:k], U[:,1:k]=myTridEigQR(SymTridiagonal(A.dv[1:k],A.ev[1:k-1]),U[:,1:k])
    λ[k+1:n], U[:,k+1:n]=myTridEigQR(SymTridiagonal(A.dv[k+1:n],A.ev[k+1:n-1]),U[:,k+1:n])
    λ, U
end

function mySymEigQR(A::Matrix)
    T,H=myTridiag(A)
    X=myTridiagX(H)
    myTridEigQR(T,X)
end

### DSTEDC
# Part of the preamble of lapack.jl
const liblapack = Base.liblapack_name
import LinearAlgebra.BLAS.@blasfunc
# import ..LinAlg: BlasFloat, Char, BlasInt, LAPACKException,
    # DimensionMismatch, SingularException, PosDefException, chkstride1, chksquare
import LinearAlgebra.BlasInt
function chklapackerror(ret::BlasInt)
    if ret == 0
        return
    elseif ret < 0
        throw(ArgumentError("invalid argument #$(-ret) to LAPACK call"))
    else # ret > 0
        throw(LAPACKException(ret))
    end
end

for (stedc, elty) in
    ((:dstedc_,:Float64),
    (:sstedc_,:Float32))
    @eval begin
        """
        COMPZ is CHARACTER*1
          = 'N':  Compute eigenvalues only.
          = 'I':  Compute eigenvectors of tridiagonal matrix also.
          = 'V':  Compute eigenvectors of original dense symmetric
                  matrix also.  On entry, Z contains the orthogonal
                  matrix used to reduce the original matrix to
                  tridiagonal form.
        """
        function stedc!(compz::Char, dv::Vector{$elty}, ev::Vector{$elty},
                Z::Array{$elty})
            n = length(dv)
            ldz=n
            if length(ev) != n - 1
                throw(DimensionMismatch("ev has length $(length(ev))
                        but needs one less than dv's length, $n)"))
            end
            w = deepcopy(dv)
            u = deepcopy(ev)
            lwork=5*n^2
            work = Array{$elty}(undef,lwork)
            liwork=6+6*n+5*n*round(Int,ceil(log(n)/log(2)))
            iwork = Array{BlasInt}(undef,liwork)
            # info = Array{BlasInt}(undef,5)
            info = Ref{BlasInt}()
            ccall((@blasfunc($stedc), liblapack), Cvoid,
                (Ref{UInt8}, Ref{BlasInt}, Ptr{$elty},
                Ptr{$elty}, Ptr{$elty}, Ref{BlasInt}, Ptr{$elty},
                    Ref{BlasInt}, Ptr{BlasInt}, Ref{BlasInt},
                    Ptr{BlasInt}), compz, n, w, u, Z, ldz, work,
                    lwork, iwork, liwork, info)
                chklapackerror(info[])
            return w,Z
        end
    end
end

function myJacobi(A1::Array{T}) where T<:Real
    A=deepcopy(A1)
    n=size(A,1)
    U=Matrix{T}(I,n,n)
    # Tolerance for rotation
    tol=sqrt(map(T,n))*eps(T)
    # Counters
    p=n*(n-1)/2
    sweep=0
    pcurrent=0
    # First criterion is for standard accuracy, second one is for relative accuracy
    # while sweep<30 && vecnorm(A-diagm(diag(A)))>tol
    while sweep<30 && pcurrent<p
        sweep+=1
        # Row-cyclic strategy
        for i = 1 : n-1
            for j = i+1 : n
                # Check for the tolerance - the first criterion is standard,
                # the second one is for relative accuracy for PD matrices
                # if A[i,j]!=zero(T)
                if abs(A[i,j])>tol*sqrt(abs(A[i,i]*A[j,j]))
                    # Compute c and s
                    τ=(A[i,i]-A[j,j])/(2*A[i,j])
                    t=sign(τ)/(abs(τ)+sqrt(1+τ^2))
                    c=1/sqrt(1+t^2)
                    s=c*t
                    G=LinearAlgebra.Givens(i,j,c,s)
                    # A=G*A
                    lmul!(G,A)
                    # A*=G'
                    rmul!(A,adjoint(G))
                    A[i,j]=zero(T)
                    A[j,i]=zero(T)
                    # U*=G'
                    rmul!(U,adjoint(G))
                    pcurrent=0
                else
                    pcurrent+=1
                end
            end
        end
    end
    diag(A), U
end

function myLanczos(A::Array{T}, x::Vector{T}, k::Int) where T
    n=size(A,1)
    X=Array{T}(undef,n,k)
    dv=Array{T}(undef,k)
    ev=Array{T}(undef,k-1)
    X[:,1]=x/norm(x)
    for i=1:k-1
        z=A*X[:,i]
        dv[i]=X[:,i]⋅z
        # Three-term recursion
        if i==1
            z=z-dv[i]*X[:,i]
        else
            # z=z-dv[i]*X[:,i]-ev[i-1]*X[:,i-1]
            # Full reorthogonalization - once or even twice
            z=z-sum([(z⋅X[:,j])*X[:,j] for j=1:i])
            z=z-sum([(z⋅X[:,j])*X[:,j] for j=1:i])
        end
        μ=norm(z)
        if μ==0
            Tr=SymTridiagonal(dv[1:i-1],ev[1:i-2])
            return eigvals(Tr), X[:,1:i-1]*eigvecs(Tr), X[:,1:i-1], μ
        else
            ev[i]=μ
            X[:,i+1]=z/μ
        end
    end
    # Last step
    z=A*X[:,end]
    dv[end]=X[:,end]⋅z
    z=z-dv[end]*X[:,end]-ev[end]*X[:,end-1]
    μ=norm(z)
    Tr=SymTridiagonal(dv,ev)
    eigvals(Tr), X*eigvecs(Tr), X, μ
end

# Extract X
function myBidiagX(H::Matrix)
    m,n=size(H)
    T=typeof(H[1,1])
    X = Matrix{T}(I,m,n)
    v = Array{T}(undef,m)
    for j = n : -1 : 1
        v[j] = one(T)
        v[j+1:m] = H[j+1:m, j]
        γ = -2 / (v[j:m]⋅v[j:m])
        w = γ * X[j:m, j:n]'*v[j:m]
        X[j:m, j:n] = X[j:m, j:n] + v[j:m]*w'
    end
    X
end

# Extract Y
function myBidiagY(H::AbstractMatrix)
    n,m=size(H)
    T=typeof(H[1,1])
    Y = Matrix{T}(I,n,n)
    v = Array{T}(undef,n)
    for j = n-2 : -1 : 1
        v[j+1] = one(T)
        v[j+2:n] = H[j+2:n, j]
        γ = -2 / (v[j+1:n]⋅v[j+1:n])
        w = γ * Y[j+1:n, j+1:n]'*v[j+1:n]
        Y[j+1:n, j+1:n] = Y[j+1:n, j+1:n] + v[j+1:n]*w'
    end
    Y
end

# Bidiagonalization using Givens rotations
function myBidiagG(A1::Matrix)
    A=deepcopy(A1)
    m,n=size(A)
    T=typeof(A[1,1])
    X=Matrix{T}(I,m,m)
    Y=Matrix{T}(I,n,n)
    for j = 1 : n
        for i = j+1 : m
            G,r=givens(A,j,i,j)
            # Use the faster in-place variant
            # A=G*A
            lmul!(G,A)
            # X=G*X
            lmul!(G,X)
        end
        for i=j+2:n
            G,r=givens(A',j+1,i,j)
            # A*=adjoint(G)
            rmul!(A,adjoint(G))
            # Y*=adjoint(G)
            rmul!(Y,adjoint(G))
        end
    end
    X',Bidiagonal(diag(A),diag(A,1),'U'), Y
end

function myJacobiR(A1::AbstractMatrix)
    A=deepcopy(A1)
    m,n=size(A)
    T=typeof(A[1,1])
    V=Matrix{T}(I,n,n)
    # Tolerance for rotation
    tol=sqrt(map(T,n))*eps(T)
    # Counters
    p=n*(n-1)/2
    sweep=0
    pcurrent=0
    # First criterion is for standard accuracy, second one is for relative accuracy
    # while sweep<30 && vecnorm(A-diagm(diag(A)))>tol
    while sweep<30 && pcurrent<p
        sweep+=1
        # Row-cyclic strategy
        for i = 1 : n-1
            for j = i+1 : n
                # Compute the 2 x 2 sumbatrix of A'*A
                F=A[:,[i,j]]'*A[:,[i,j]]
                # Check the tolerance - the first criterion is standard,
                # the second one is for relative accuracy
                # if A[i,j]!=zero(T)
                #
                if abs(F[1,2])>tol*sqrt(F[1,1]*F[2,2])
                    # Compute c and s
                    τ=(F[1,1]-F[2,2])/(2*F[1,2])
                    t=sign(τ)/(abs(τ)+sqrt(1+τ^2))
                    c=1/sqrt(1+t^2)
                    s=c*t
                    G=LinearAlgebra.Givens(i,j,c,s)
                    A*=G'
                    # rmul!(A,adjoint(G))
                    V*=G'
                    # rmul!(A,adjoint(G))
                    pcurrent=0
                else
                    pcurrent+=1
                end
            end
        end
    end
    σ=[norm(A[:,k]) for k=1:n]
    for k=1:n
        A[:,k]./=σ[k]
    end
    A, σ, V
end

end

src/Module B - Eigenvalue and Singular Value Decompositions/ModuleB.jl


import FractionalTimeDG.P, FractionalTimeDG.dP

α = 3/4
r = 4
store = FractionalTimeDG.Store(α, r, 2r)

function brute_force_H0(r::Integer, α::T, rtol=1e-8) where T <: AbstractFloat

    # \int_0^\tau (\tau-\sigma)^{\alpha-1} P_{j-1}(\tau)\,d\tau
    function inner_integral(j::Integer, τ::T) where T <: AbstractFloat
        I, err = quadgk(-one(T), τ, rtol=rtol) do σ
            (τ-σ)^(α-1) * P(j-1, σ)
        end
        return I
    end

    H0 = Array{T}(undef, r, r)
    c = 1 / (2^α*Γ(α))
    for j = 1:r, i = 1:r
        I, err = quadgk(-one(T), one(T), rtol=rtol) do τ
            dP(i-1, τ) * inner_integral(j, τ)
        end
        H0[i,j] = c * ( inner_integral(j, one(T) ) - I )
    end
    return H0
end

function brute_force_H_uniform(ℓbar::Integer, r::Integer, α::T,
                       rtol=1e-8) where T <: AbstractFloat

    function inner_integral(ℓbar::Integer, j::Integer, τ::T
                           ) where T <: AbstractFloat
        I, err = quadgk(-one(T), one(T), rtol=rtol) do σ
            Δ = (τ-σ) / ( 2ℓbar )
            return (1+Δ)^(α-1) * P(j-1, σ)
        end
        return I
    end

    Hℓbar = Array{T}(undef, r, r)
    c = ℓbar^(α-1) / ( 2Γ(α) )
    for j = 1:r
        Aj = inner_integral(ℓbar, j, one(T))
        Bj = inner_integral(ℓbar, j, -one(T))
        pow = one(T)
        for i = 1:r
            Cij, err = quadgk(-one(T), one(T), rtol=rtol) do τ
                dP(i-1,τ) * inner_integral(ℓbar, j, τ)
            end
            pow = -pow
            Hℓbar[i,j] = c * ( Aj + pow * Bj - Cij )
        end
    end
    return Hℓbar
end

N = 4

H = coef_H_uniform!(r, N, r+2, store)

bf_H0 = brute_force_H0(r, α)

err0 = bf_H0 - H[0]
@test all( abs.(err0) .< 1e-8 )

bf_H2 = brute_force_H_uniform(2, r, α)
err2 = bf_H2- H[2]
@test all( abs.(err2) .< 1e-8 )

bf_H3 = brute_force_H_uniform(3, r, α)
err3 = bf_H3- H[3]
@test all( abs.(err3) .< 1e-8 )


using VectorSphericalWaves

m, n, kr_r, kr_i, θ, ϕ = 1, 1, 1e6, 1e2, 0.2, 0.3;
kr = complex(kr_r, kr_i);

M_mn_wave(m, n, kr, θ, ϕ; kind="regular");
N_mn_wave(m, n, kr, θ, ϕ; kind="regular");

# BigFloat
m, n, kr_r, kr_i, θ, ϕ = 1, 1, BigFloat(1e6), BigFloat(1e2), BigFloat(0.2), BigFloat(0.3);
kr = complex(kr_r, kr_i);

M_mn_wave(m, n, kr, θ, ϕ; kind="regular");
N_mn_wave(m, n, kr, θ, ϕ; kind="regular");

# SVector comparison
m, n, kr_r, kr_i, θ, ϕ = 1, 1, 1e6, 1e2, 0.2, 0.3;
kr = complex(kr_r, kr_i);

θ_array, ϕ_array = LinRange(0, π, 1000), LinRange(0, π, 1000)

println("M and N returning Arrays")
@time M_mn_wave.(m, n, kr, θ_array, ϕ_array; kind="regular");
@time N_mn_wave.(m, n, kr, θ_array, ϕ_array; kind="regular");

println("M and N returning SVector, you should notice the sligth speedup")
@time VectorSphericalWaves.M_mn_wave_SVector.(m, n, kr, θ_array, ϕ_array; kind="regular");
@time VectorSphericalWaves.N_mn_wave_SVector.(m, n, kr, θ_array, ϕ_array; kind="regular");


using LinearAlgebra
using Combinatorics

A = [3 1 2 2;
     1 3 8 2;
     2 3 5 1]

B = [1 0 2;
     1 3 6;
     2 2 0;
     0 1 1]

function caychy_binet(A, B)
    comb = collect(combinations(1:size(A)[2], size(A)[1]))
    res = 0
    for S in comb
        res += det(A[:,S])*det(B[S,:])
    end

    res
end

println("det(AB)=", det(A*B))
println("cauchy-binet: ", caychy_binet(A, B))


#
# derivative()
#

f1(x::Real) = sin(x)
@test norm(derivative(f1, :scalar, :forward)(0.0) - cos(0.0)) < 10e-4
@test norm(derivative(f1, :scalar, :central)(0.0) - cos(0.0)) < 10e-4
@test norm(derivative(f1, :forward)(0.0) - cos(0.0)) < 10e-4
@test norm(derivative(f1, :central)(0.0) - cos(0.0)) < 10e-4
@test norm(derivative(f1)(0.0) - cos(0.0)) < 10e-4

f2(x::Vector) = sin(x[1])
@test norm(derivative(f2, :vector, :forward)([0.0]) .- cos(0.0)) < 10e-4
@test norm(derivative(f2, :vector, :central)([0.0]) .- cos(0.0)) < 10e-4

#
# adjoint overloading
#

f3(x::Real) = sin(x)
for x in range(0.0, stop=0.1, length=11) # seq()
    @test norm(derivative(f3,x) - cos(x)) < 10e-4
end

#
# gradient()
#

f4(x::Vector) = (100.0 - x[1])^2 + (50.0 - x[2])^2
@test norm(Calculus.gradient(f4, :forward)([100.0, 50.0]) - [0.0, 0.0]) < 10e-4
@test norm(Calculus.gradient(f4, :central)([100.0, 50.0]) - [0.0, 0.0]) < 10e-4
@test norm(Calculus.gradient(f4)([100.0, 50.0]) - [0.0, 0.0]) < 10e-4

#
# jacobian()
#

@test norm(Calculus.jacobian(identity, rand(3), :forward) - Matrix(1.0I, 3, 3)) < 10e-4
@test norm(Calculus.jacobian(identity, rand(3), :central) - Matrix(1.0I, 3, 3)) < 10e-4

#
# second_derivative()
#

@test norm(second_derivative(x -> x^2)(0.0) - 2.0) < 10e-4
@test norm(second_derivative(x -> x^2)(1.0) - 2.0) < 10e-4
@test norm(second_derivative(x -> x^2)(10.0) - 2.0) < 10e-4
@test norm(second_derivative(x -> x^2)(100.0) - 2.0) < 10e-4

#
# hessian()
#

f5(x) = sin(x[1]) + cos(x[2])
@test norm(Calculus.gradient(f5)([0.0, 0.0]) - [cos(0.0), -sin(0.0)]) < 10e-4
@test norm(hessian(f5)([0.0, 0.0]) - [-sin(0.0) 0.0; 0.0 -cos(0.0)]) < 10e-4


using PyPlot, LinearAlgebra, Roots, DifferentialEquations

function f(u,v; b=0.10, i=0.00)
	return u*(1.0-u)*(u-b) - v + i
end

function g(u,v; a=0.37, c=0.05)
	return c*(a*u-v)
end

function FHNgrid(; a=0.37, b=0.1, c=0.05, i=0.0, NU=1025, NV=1025, ulims=[-2.,+2.],vlims=[-1.,+1.])
	uu = range(ulims[1],ulims[2],NU)
	vv = range(vlims[1],vlims[2],NV)
	U = zeros(Float64, NV, NU)
	V = zeros(Float64, NV, NU)
	F = zeros(Float64, NV, NU)
	G = zeros(Float64, NV, NU)

	for m in 1:NV, n in 1:NU
		U[m,n] = uu[n]
		V[m,n] = vv[m]
		F[m,n] = f(U[m,n],V[m,n];b=b,i=i)
		G[m,n] = g(U[m,n],V[m,n];a=a,c=c)
	end

	return U,V,F,G
end

function FHNnullclines(; a=0.37, b=0.1, c=0.05, i=0.0)

	U,V,F,G = FHNgrid(; a=a, b=b, c=c, i=i)
	fig = plt.contour(U,V,F,levels=[-1e-14,+1e-14],colors="C0",linestyles="solid")
	plt.contour(U,V,G,levels=[-1e-14,+1e-14],colors="C1",linestyles="solid")
	plt.xlabel("\$ u \$")
	plt.ylabel("\$ v \$")
	plt.xlim([-2,2])
	plt.ylim([-1,1])
	return fig
end

function FHNequilibria(; a=0.37, b=0.1, c=0.05, i=0.0)

	# need to find ALL u s.t.:
	#	f(u,v) = 0 && g(u,v) = 0
	#	Since g(u,v) = 0 is linear, we invert to get v = a*u
	# plugging this into f(u,v=a*u) = 0, we have a single
	# scalar, cubic, equation to find the roots of.
	# There will be between 0 and 3 roots.
	ulim = 2.0
	r = find_zeros((u)->f(u,a*u;b=b,i=i),-ulim,ulim)
	while length(r) < 1 && ulim < 100.0
		ulim = 2*ulim
		r = find_zeros((u)->f(u,a*u;b=b,i=i),-ulim,ulim)
	end
	return reduce(hcat, (r,a*r))

	# NOTE: for the general case, where g(u,v) = 0 is not invertible,
	# you must use Newton's method:
	#	[u;v] = [u;v] - FHNjac(u,v)\[f(u,v);g(u,v)]
	# in a loop, until convergence (u & v do not change on successive iterations)
	# AND this only works for one root at a time, whereas we need all three!
	# AND we need good estimates to initialize this Newton iteration!
	#u = randn(); v = randn()
	#for n in 1:20
	#	(u,v) = (u,v) .- FHNjac(u,v; a=a,b=b,c=c,i=i)\[f(u,v;b=b,i=i);g(u,v;a=a,c=c)]
	#end
end

function FHNjac(u,v; a=0.37, b=0.1, c=0.05, i=0.0, h=1e-14)
	# this uses Complex-Step differentiation:
	# 	https://blogs.mathworks.com/cleve/2013/10/14/complex-step-differentiation/
	# a method of computing derivatives efficiently and accurately
	# using complex arithmetic.
	# It may, however, fail if the functions are not holonomic!
	J = zeros(Float64,2,2)
	J[1,1] = imag(f(u+im*h,v;b=b,i=i))/h
	J[1,2] = imag(f(u,v+im*h;b=b,i=i))/h
	J[2,1] = imag(g(u+im*h,v;a=a,c=c))/h
	J[2,2] = imag(g(u,v+im*h;a=a,c=c))/h

	return J
end

function stability(J)
	return maximum(real(eigvals(J)))
end

function stabilityplot(;N=10000)

	trJ = zeros(Float64,N)
	detJ = zeros(Float64, N)
	lJ = zeros(Float64, N)
	for n in 1:N
		J = randn(Float64, 2, 2)
		lJ[n] = stability(J)
		trJ[n] = tr(J)
		detJ[n] = det(J)
	end
	plt.scatter(trJ, detJ, c=lJ, s=1.0, cmap="bwr", vmin=-1, vmax=+1)
	plt.colorbar(label="\$ \\max \\, \\Re (\\lambda) \$", extend="both")
	plt.plot(-5.0:0.1:5.0, ((-5.0:0.1:5.0).^2.0)./4.0, "-k", label="\$ \\Delta = 0 \$")
	plt.xlabel("tr\$ A \$")
	plt.ylabel("\$ \\det A \$")
	plt.xlim([-5,5])
	plt.ylim([-5,5])
	plt.legend(loc=0, edgecolor="none")
end

function oscillationplot(;N=10000)

	trJ = zeros(Float64,N)
	detJ = zeros(Float64, N)
	lJ = zeros(Float64, N)
	for n in 1:N
		J = randn(Float64, 2, 2)
		lJ[n] = maximum(abs.(imag.(eigvals(J))))
		trJ[n] = tr(J)
		detJ[n] = det(J)
	end
	plt.scatter(trJ, detJ, c=lJ, s=1.0, cmap="bwr", vmin=-1, vmax=+1)
	plt.colorbar(label="\$ \\max \\, \\Im (\\lambda) \$", extend="both")
	plt.plot(-5.0:0.1:5.0, ((-5.0:0.1:5.0).^2.0)./4.0, "-k", label="\$ \\Delta = 0 \$")
	plt.xlabel("tr\$ A \$")
	plt.ylabel("\$ \\det A \$")
	plt.xlim([-5,5])
	plt.ylim([-5,5])
	plt.legend(loc=0, edgecolor="none")
end

function FHNstrmplt(; a=0.37, b=0.1, c=0.05, i=0.0)
	U,V,F,G = FHNgrid(; a=a, b=b, c=c, i=i)

	fig = plt.streamplot(U,V,F,G, color=sqrt.(F.^2 + G.^2))
	plt.colorbar(label="\$ \\sqrt{\\dot u ^2 + \\dot v ^2} \$")
	plt.xlabel("\$ u \$")
	plt.ylabel("\$ v \$")
	plt.xlim([-2,2])
	plt.ylim([-1,1])
	return fig
end

function FHNtotal(; a=0.37, b=0.1, c=0.05, i=0.0)
	fig = plt.figure()
	FHNstrmplt(;a=a,b=b,c=c,i=i)
	FHNnullclines(;a=a,b=b,c=c,i=i)
	r = FHNequilibria(;a=a,b=b,c=c,i=i)
	for n in 1:size(r,1)
		u = r[n,1]; v = r[n,2]
		l = eigvals(FHNjac(u,v;a=a,b=b,c=c,i=i))
		ii = argmax(real(l))
		if maximum(real(l)) > 0
			plt.plot([u],[v],"ok", markerfacecolor="w", label="\$ \\lambda = $(round(l[ii],sigdigits=2)) \$")
		elseif maximum(real(l)) < 0
			plt.plot([u],[v],"ok", label="\$ \\lambda = $(round(l[ii],sigdigits=2)) \$")
		end
	end
	plt.legend(loc=0, edgecolor="none")
	return fig
end

function FHNdyn(; a=0.37, b=0.1, c=0.05, i=0.0, x0=randn(Float64,2))
	function fhn!(dx, x, p, t)
		dx[1] = f(x[1],x[2];b=p[2],i=p[4])
		dx[2] = g(x[1],x[2];a=p[1],c=p[3])
		return nothing
	end
	prob = ODEProblem(fhn!, x0, (0.0, 1000.0), [a,b,c,i])
	sol = solve(prob, abstol=1e-13, reltol=1e-13)

	FHNtotal(;a=a,b=b,c=c,i=i)
	plt.plot(sol[1,:], sol[2,:], "-k", linewidth=3)
	plt.plot(x0[1],x0[2],"xk")
	plt.title("\$ \\alpha=$(a), \\beta=$(b), \\gamma=$(c), I=$(round(i,sigdigits=3)) \$")
	plt.xlabel("\$ u \$")
	plt.ylabel("\$ v \$")
	plt.xlim([-.5,1.5])
	plt.ylim([-.5,.5])
	return nothing

end

function FHNexc(; a=0.37, b=0.1, c=0.05, i=0.0)

	U,V,F,G = FHNgrid(; a=a, b=b, c=c, i=i, ulims=[-0.3,1.3], vlims=[-0.3,0.3], NU=129, NV=65)

	fig, axs = plt.subplots(3,4,sharex=true, sharey=true, constrained_layout=true)
	for (n,u) in enumerate(range(-0.1,0.2,step=0.1)), (m,v) in enumerate(range(0.05,-0.05,step=-0.05))
		axs[m,n].streamplot(U,V,F,G,color=sqrt.(F.^2 + G.^2),linewidth=sqrt.(F.^2 + G.^2))
		axs[m,n].contour(U,V,F,levels=[-1e-14,+1e-14],colors="C0",linestyles="solid")
		axs[m,n].contour(U,V,G,levels=[-1e-14,+1e-14],colors="C1",linestyles="solid")

		function fhn!(dx, x, p, t)
			dx[1] = f(x[1],x[2];b=p[2],i=p[4])
			dx[2] = g(x[1],x[2];a=p[1],c=p[3])
			return nothing
		end
		prob = ODEProblem(fhn!, [u;v], (0.0, 1000.0), [a,b,c,i])
		sol = solve(prob, abstol=1e-13, reltol=1e-13)

		axs[m,n].plot(sol[1,:], sol[2,:], "-k", linewidth=2)
		axs[m,n].plot(u,v,"xk")
		axs[m,n].plot(0,0,".k")

		if n==1
			axs[m,n].set_ylabel("\$ v_0 = $(v) \$")
		end
		if m==3
			axs[m,n].set_xlabel("\$ u_0 = $(u) \$")
		end
		axs[m,n].set_xlim([-0.3,1.3])
		axs[m,n].set_ylim([-0.3,0.3])
	end

	return fig, axs

end

function currentBif(;a=0.37,b=0.1,c=0.05,i=0.0)
	ii = range(0.00,0.05,step=0.0001)
	li = zeros(Float64, length(ii))
	for (n,i) in enumerate(ii)
		x0 = FHNequilibria(;i=i)
		x0 = x0[argmin(x0[:,1]),:]
		li[n] = stability(FHNjac(x0[1],x0[2];i=i))
	end
	fig = figure()
	plt.plot(ii, li, "-", label="\$ \\max \\Re( \\lambda)(I) \$")
	ind = argmin(abs.(li))
	plt.plot(ii[ind], li[ind], ".k", label="\$ I^* = $(ii[ind]) \$")
	plt.xlabel("\$ I \$")
	plt.ylabel("")
	plt.legend(loc=0,edgecolor="none")
	return fig
end

function currentBifTraj(;a=0.37,b=0.1,c=0.05,i=0.0)
	function fhn!(dx, x, p, t)
		dx[1] = f(x[1],x[2];b=p[2],i=p[4])
		dx[2] = g(x[1],x[2];a=p[1],c=p[3])
		return nothing
	end
	fig, axs = plt.subplots(5,1,sharex=true,sharey=true,constrained_layout=true)

	for (n,i) in enumerate(range(0.0,0.04,length=5))
		prob = ODEProblem(fhn!, [1.0;-0.1], (0.0, 1000.0), [a,b,c,i])
		sol = solve(prob, abstol=1e-13, reltol=1e-13)
		axs[n].plot(sol.t, sol[1,:], "-C0", label="\$ u(t) \$")
		axs[n].plot(sol.t, sol[2,:], "-C1", label="\$ v(t) \$")
		axs[n].set_ylabel("\$ I = $(round(i,sigdigits=2)) \$")
	end
	axs[1].legend(loc=0, edgecolor="none")
	axs[end].set_xlabel("\$ t \$")
	axs[end].set_xlim([0,200])
	plt.savefig("/Users/christophermarcotte/Downloads/fhncurrentbiftraj.pdf")
	plt.close("all")
end

function alphaBifPlot(;a=0.37,b=0.1,c=0.05,i=0.0)
	fig = plt.figure()
	astar = 0.2025
	for a in range(0.37, -0.07, length=4096)
		rr = FHNequilibria(;a=a,b=b,c=c,i=i)
		for n in 1:size(rr,1)
			J = FHNjac(rr[n,1],rr[n,2]; a=a,b=b,c=c,i=i)
			l = eigvals(J)
			if maximum(real(l)) > 0
				plt.plot(a,rr[n,1],".k")
			elseif maximum(real(l)) < 0
				plt.plot(a, rr[n,1],"ok")
			end
		end
	end
	plt.xlabel("\$ \\alpha \$")
	plt.ylabel("\$ \\bar{u} \$")
	plt.title("Thick - Stable; Thin - Unstable")
	return fig
end

function FHNLC(; a=0.37, b=0.1, c=0.05, i=0.04)
	FHNtotal(;a=a,b=b,c=c,i=i)
	function fhn!(dx, x, p, t)
		dx[1] = f(x[1],x[2];b=p[2],i=p[4])
		dx[2] = g(x[1],x[2];a=p[1],c=p[3])
		return nothing
	end
	prob = ODEProblem(fhn!, [0.0;-0.1], (0.0, 1000.0), [a,b,c,i])
	sol = solve(prob, abstol=1e-13, reltol=1e-13)
	plt.plot(sol[1,:], sol[2,:], "-k", linewidth=2)
	plt.plot(sol[1,1], sol[2,1],"xk")

	prob = ODEProblem(fhn!, [0.0;0.1], (0.0, 1000.0), [a,b,c,i])
	sol = solve(prob, abstol=1e-13, reltol=1e-13)
	plt.plot(sol[1,:], sol[2,:], "-k", linewidth=2)
	plt.plot(sol[1,1], sol[2,1],"xk")

	plt.title("\$ \\alpha=$(a), \\beta=$(b), \\gamma=$(c), I=$(round(i,sigdigits=3)) \$")
	plt.xlabel("\$ u \$")
	plt.ylabel("\$ v \$")
	plt.xlim([-.5,1.5])
	plt.ylim([-.2,.3])

	plt.savefig("/Users/christophermarcotte/Downloads/fhnLC.pdf")
	plt.close("all")
end


# van_der_Pol_Poincare_section_v1.jl
# =============================================

using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, OrdinaryDiffEq
using StaticArrays
using Random
using Statistics

function van_der_Pol_ODE(du,u,p,t)
    d, a, omega = p
    du[1] =  u[2]
    du[2] =  - u[1] - d*(u[1]*u[1]-1) * u[2] + a*sin(omega*t)
end

# parameters and initial values
# -----------------------------
ntrans = 5000.   # transient periods
nattr = 2000    # no. of points on the attractor used for estimating winding no.
y_init = [0.1 0.1]

d = 5.
a = 1.
omega = 1.

p = [d a omega]
period = 2*pi / omega

 # transient to attractor
 # ----------------------
tspan = (0., ntrans*period)
prob = ODEProblem(van_der_Pol_ODE, y_init, tspan, p)
sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8, maxiters = 1e12 )

# attractor
# ---------
tspan = (0., nattr*period)
y_init = sol.u[end]
prob = ODEProblem(van_der_Pol_ODE, y_init, tspan, p)
sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8, maxiters = 1e12,  saveat = period)

u1vec = [u[1] for u in sol.u]
u2vec = [u[2] for u in sol.u]
nts = length(u1vec)

theta_vec = (atan.( u2vec , u1vec  ) .+ pi ) ./ (2*pi)

# open plot
# ---------
# %%
fig = figure(figsize=(0.6*figx,figy))

ax1 = subplot(1,2,1)
plt.scatter(u1vec,u2vec,s=4, c = COLORS[1])

xticks([-2,0,2], ["-2","","2"])
yticks([-6,0,6], ["-6","","6"])
xlabel(L"u_1";labelpad = -20)
ylabel(L"u_2";labelpad = -20)

ax2 = subplot(1,2,2)
plt.scatter(theta_vec[1:end-1],theta_vec[2:end],s=4, c = COLORS[2])
xlabel(L"\theta_n"; labelpad = -20)
ylabel(L"\theta_{n+1}"; labelpad = -20)
xticks([0,0.5,1], ["0", "", "2π"])
yticks([0,0.5,1], ["0", "", "2π"])
yticks([0,0.5,1])
xlim(0,1)
ylim(0,1)
add_identifiers!(fig)
fig.tight_layout(pad=0.33)

wsave(plotsdir("9", "vanderPol_poincare_section"), fig)


# 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 dynamic programming 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%20dynamic%20programming.jl

# In[2]:

#compute the maximum profit for long only strategy
function compute_profit(prices)

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

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

    #initialize indices
    left=1
    right=length(prices)
    minind=1
    maxind=length(prices)

    #we have two indices moving at the same time
    #one from left to find the minimum value
    #the other from right to find the maximum value
    #we are not looking for global minimum
    #we only want local minimum before the global maximum
    while left<maxind && right>minind

        #when a larger value is found
        #update accordingly
        if prices[right]>max_price

            max_price=prices[right]
            maxind=copy(right)

        end

        #the same applies to a smaller value
        if prices[left]<min_price

            min_price=prices[left]
            minind=copy(left)

        end

        left+=1
        right-=1

    end

    #maximum profit
    profit=max_price-min_price

    #when we lose money
    #abort the trade
    if profit<0

        profit=0

    end

    return profit

end

# In[3]:

#there are two scenarios to maximize the profit
#one trade or two trades
#since we can execute two transactions
#we split the array at an iterating index i into half
#we find out the maximum profit we can obtain from both halves
function stock_trading(prices,ind)

    #prevent index error
    if ind+1==length(prices)

        return 0

    end

    #split
    upper=prices[1:ind]
    lower=prices[ind+1:end]

    #compute profit
    profit=compute_profit(lower)+compute_profit(upper)

    #compare recursively
    return max(profit,stock_trading(prices,ind+1))

end

# In[4]:

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

# In[5]:

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

# In[6]:

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

# In[7]:

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