"""
1-D Gradient

    gradient(a::Array, dx::Float64=1)
    :param f: Vector
    :param dx: stepsize

    :return g: gradient
"""
function gradient(f::Array{Float64, 1}, dx::Float64=1.)
    n = length(f)
    g = zeros(n)
    h = dx*(1:n)
    # Take forward differences on left and right edges
    g[1] = (f[2] - f[1])/(h[2]-h[1])
    g[n] = (f[n] - f[n-1])/(h[end]-h[end-1])

    # Take centered differences on interior points
    h = h[3:n]-h[1:(n-2)]
    g[2:(n-1)] = (f[3:n]-f[1:(n-2)])/h[1]

    return g
end

"""
2-D Gradient

    gradient2(a::Array, dx::Float64=1, dy::Float=1)
    :param a: matrix
    :param dx: stepsize
    :param dy: stepsize

    :return dxdu: derivatives along x
    :return dydv: derivatives along y
"""
function gradient2(a::Array, dx::Float64=1., dy::Float64=1.)
    m,n = size(a);
    dxdu = zeros(m,n);
    dydv = zeros(m,n);

    for i=1:m
        dxdu[i,:] = gradient(vec(a[i,:]), dx)
    end

    for i=1:n
        dydv[:,i] = gradient(a[:,i], dy)
    end

    return dxdu, dydv
end

"""
Creates Rectangular Grid in 2-D space

    meshgrid(a::LinRange,b::LinRange)
"""
function meshgrid(a::LinRange,b::LinRange)
    grid_a = [i for i in a, j in b]';
    grid_b = [j for i in a, j in b]';

    return grid_a, grid_b
end

"""
Linear interpolation

    approx(xd, yd, xi)
    :param xd: x samples
    :param yd: response samples
    :param xi: new x samples
"""
function approx(xd, yd, xi)
    nd = length(xd);
    ni = length(xi);

    yi = zeros(ni);
    for i in 1:ni
        if (xi[i] <= xd[1])
            t = (xi[i]-xd[1]) / (xd[2] - xd[1]);
            yi[i] = (1.0 - t) * yd[1] + t * yd[2];
        elseif (xd[nd] <= xi[i])
            t = (xi[i] - xd[nd-1]) / (xd[nd] - xd[nd-1]);
            yi[i] = (1.0 - t) * yd[nd-1] + t * yd[nd];
        else
            for k in 2:nd
                if (xd[k-1] <= xi[i] && xi[i] <= xd[k])
                    t = (xi[i] - xd[k-1]) / (xd[k] - xd[k-1]);
                    yi[i] = (1.0 - t) * yd[k-1] + t * yd[k];
                    break
                end
            end
        end
    end

    return yi
end

"""
Trapezoidal Integration

    trapz(x, y, dim=1)
    :param x: vector of time samples
    :param y: array of response samples
    :param dim: dimension along which to integrate
"""
function trapz(x::Array{Float64, 1}, y::Array{Float64}, dim::Integer=1)
    perm = [dim:max(ndims(y),dim); 1:dim-1];
    y = permutedims(y, perm);
    if ndims(y) == 1
        m = 1;
    else
        m = size(y,1);
    end

    if m == 1
        M = length(y);
        out = sum(diff(x).*(y[1:M-1] + y[2:M])/2.0);
    else
        out = transpose(diff(x)) * (reshape(y,Val(2))[1:m-1,:] + reshape(y,Val(2))[2:m,:])/2.0;
        siz = size(y); siz = collect(siz); siz[1] = 1;
        out = reshape(out, tuple(siz...));
        out = permutedims(out, invperm(perm));
        ind = findall(collect(size(out)).==1);
        out = dropdims(out,dims=ind[1]);
        if length(out) == 1;
            out = out[1];
        end
    end

    return out
end

"""
Cumulative Trapezoidal Integration

    cumtrapz(x, y, dim=1)
    :param x: vector describing time samples
    :param y: array describing response
    :param dim: dimension to integrate over
"""
function cumtrapz(x::Array{Float64, 1}, y::Array{Float64}, dim::Integer=1)
    perm = [dim:max(length(size(y)),dim); 1:dim-1];
    y = permutedims(y, perm);
    if ndims(y) == 1
        n = 1;
        m = length(y);
    else
        m, n = size(y);
    end

    if n == 1
        dt = diff(x)/2.0;
        z = [0; cumsum(dt.*(y[1:(m-1)] + y[2:m]))];
    else
        dt = repeat(diff(x)/2.0,1,n);
        z = [zeros(1,n); cumsum(dt.*(y[1:(m-1), :] + y[2:m, :]),dims=1)];
        z = permutedims(z, invperm(perm));
    end

    return z
end

"""
Cumulative Trapezoidal Integration using midpoint

    cumtrapzmid(x, y, c)
    :param x: time samples
    :param y: resposne samples
    :param c: midpoint
    :param mid: midpoint location
"""
function cumtrapzmid(x, y, c, mid)
    a = length(x);

    # case < mid
    fn = zeros(a);
    tmpx = x[(mid-1):-1:1];
    tmpy = y[(mid-1):-1:1];
    tmp = c .+ cumtrapz(tmpx, tmpy);
    fn[1:(mid-1)] = reverse(tmp);

    # case >= mid
    fn[mid:a] = c .+ cumtrapz(x[mid:a],y[mid:a]);

    return fn

end

"""
Multivariate Normal random number generation

    mvnrand(mu, C, n)
    :param mu: mean vector
    :param C: covariance matrix
    :param n: number of samples
"""
function mvnrand(mu, C, n)
    tmp = cholesky(C, :U, Val{true});
    R = tmp[:U];
    R = Array(R);
    R = R[:, tmp.piv];
    retval = randn(n, size(R,1)) * R;
    retval += transpose(repeat(mu, 1, n));
    return transpose(retval)
end

"""
Linear interpolation when response contains flat regions

    interp1_flat(x, y, xx)
    :param x: time samples
    :param y: response samples
    :param xx: new time samples
"""
function interp1_flat(x,y,xx)
    flat = find(diff(x).<=0);
    n = length(flat);

    if n==0
        tmp = interpolate((x,), y, Gridded(Linear()))
        yy = tmp[xx];
    else
        yy = zeros(size(xx));
        i1 = 1;
        if flat[1] == 1
            i2 = 1;
            j = xx.==x[i2];
            yy[j] = minimum(y[i2:i2+1]);
        else
            i2 = flat[1];
            j = (xx.>=x[i1]) & (xx.<=x[i2]);
            tmp = interpolate((x[i1:i2],), y[i1:i2], Gridded(Linear()))
            yy[j] = tmp[xx[j]];
            i1 = copy(i2);
        end
        for k = 2:n
            i2 = flat[k];
            if i2 > i1+1
                j = (xx.>=x[i1]) & (xx.<=x[i2]);
                yi = interpolate((x[i1+1:i2],), y[i1+1:i2], Gridded(Linear()))
                yy[j] = tmp[xx[j]];
            end
            j = xx.==x[i2];
            yy[j] = minimum(y[i2:i2+1]);
            i1 = copy(i2);
        end
        i2 = length(x);
        j = (xx.>=x[i1]) & (xx.<=x[i2]);
        if i1+1 == i2
            yy[j] = y[i2];
        else
            tmp = interpolate((x[i1+1:i2],), y[i1+1:i2], Gridded(Linear()))
            yy[j] = tmp[xx[j]];
        end
    end

    return yy
end

function basis_fourier(f_domain::Vector, numBasis::Int64, fourier_p::Int64)
    result = Array{Float64}(undef,length(f_domain),2*numBasis)
    for i in 1:2*numBasis
        j = ceil(i/2)
        if ((i % 2) == 1)
            result[:,i] = sqrt(2) * sin.(2*j*pi*f_domain/fourier_p)
        end
        if ((i % 2) == 0)
            result[:,i] = sqrt(2) * cos.(2*j*pi*f_domain/fourier_p)
        end
    end

    out = basis(f_domain,result)
    return out
end

function f_predictfunction(f::func, at; deriv=0, method="linear")
    if method=="linear"
        if (deriv==0)
            result = approx(f.x, f.y, at)
        end
        if (deriv == 1)
            fmod = approx(f.x, f.y, at)
            diffy1 = [0; diff(fmod)]
            diffy2 = [diff(fmod); 0]
            diffx1 = [0; diff(at)]
            diffx2 = [diff(at); 0]

            result = (diffy2 + diffy1) ./ (diffx2 + diffx1)
        end
    else
        error("Method not implemented")
    end

    out = func(at, result)
    return out
end

function f_basistofunction(f_domain, coef, cbasis::basis; coefconst=0)
    if (size(cbasis.matrix,2)<length(coef))
        error("coefficients exceeds basis functions")
    end

    y1 = cbasis.matrix[:,(1:length(coef))] * coef .+ coefconst
    result = func(cbasis.x, y1)
    y1 = f_predictfunction(result, f_domain)

    return y1
end


# elastic-net regularization

export ElasticNet

"""
**Elastic-net regularization**

    ElasticNet(μ=1.0, λ=1.0)

Returns the function
```math
f(x) = μ\\|x\\|_1 + (λ/2)\\|x\\|^2,
```
for nonnegative parameters `μ` and `λ`.
"""
struct ElasticNet{R <: Real} <: ProximableFunction
    mu::R
    lambda::R
    function ElasticNet{R}(mu::R, lambda::R) where {R <: Real}
        if lambda < 0 || mu < 0
            error("parameters `μ` and `λ` must be nonnegative")
        else
            new(mu, lambda)
        end
    end
end

is_separable(f::ElasticNet) = true
is_prox_accurate(f::ElasticNet) = true
is_convex(f::ElasticNet) = true

ElasticNet(mu::R=1.0, lambda::R=1.0) where {R <: Real} = ElasticNet{R}(mu, lambda)

function (f::ElasticNet)(x::AbstractArray{R}) where R <: RealOrComplex
    return f.mu*norm(x,1) + (f.lambda/2)*norm(x,2)^2
end

function prox!(y::AbstractArray{R}, f::ElasticNet{R}, x::AbstractArray{R}, gamma::R=R(1)) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    gm = gamma*f.mu
    gl = gamma*f.lambda
    for i in eachindex(x)
        y[i] = (x[i] + (x[i] <= -gm ? gm : (x[i] >= gm ? -gm : -x[i])))/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{R}, f::ElasticNet{R}, x::AbstractArray{R}, gamma::AbstractArray{R}) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    for i in eachindex(x)
        gm = gamma[i]*f.mu
        gl = gamma[i]*f.lambda
        y[i] = (x[i] + (x[i] <= -gm ? gm : (x[i] >= gm ? -gm : -x[i])))/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{Complex{R}}, f::ElasticNet{R}, x::AbstractArray{Complex{R}}, gamma::R=R(1)) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    gm = gamma*f.mu
    gl = gamma*f.lambda
    for i in eachindex(x)
        y[i] = sign(x[i])*max(0, abs(x[i]) - gm)/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{Complex{R}}, f::ElasticNet{R}, x::AbstractArray{Complex{R}}, gamma::AbstractArray{R}) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    for i in eachindex(x)
        gm = gamma[i]*f.mu
        gl = gamma[i]*f.lambda
        y[i] = sign(x[i])*max(0, abs(x[i]) - gm)/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function gradient!(y::AbstractArray{T}, f::ElasticNet{R}, x::AbstractArray{T}) where {T <: RealOrComplex, R <: Real}
    # Gradient of 1 norm
    y .= f.mu.*sign.(x)
    # Gradient of 2 norm
    y .+= f.lambda.*x
    return f.mu*norm(x,1) + (f.lambda/2)*norm(x,2)^2
end

fun_name(f::ElasticNet) = "elastic-net regularization"
fun_dom(f::ElasticNet) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::ElasticNet) = "x ↦ μ||x||_1 + (λ/2)||x||²"
fun_params(f::ElasticNet) = "μ = $(f.mu), λ = $(f.lambda)"

function prox_naive(f::ElasticNet, x::AbstractArray{R}, gamma::Real=1.0) where R <: RealOrComplex
    uz = max.(0, abs.(x) .- gamma*f.mu)/(1 + f.lambda * gamma)
    return sign.(x) .* uz, f.mu * norm(uz,1) + (f.lambda/2) * norm(uz)^2
end

function prox_naive(f::ElasticNet, x::AbstractArray{R}, gamma::AbstractArray) where R <: RealOrComplex
    uz = max.(0, abs.(x) .- gamma.*f.mu)./(1 .+ f.lambda .* gamma)
    return sign.(x) .* uz, f.mu * norm(uz,1) + (f.lambda/2) * norm(uz)^2
end


# This is a Julia adaptation of the Matlab code from O. Ledoit and M. Wolf
# solely used to test our own implementation of their estimator.
# Their code was released under the BSD-2 license with the copyright notice
# below.
# https://www.econ.uzh.ch/en/people/faculty/wolf/publications.html#9

###########################################################################
# This file is released under the BSD 2-clause license.

# Copyright (c) 2014, Olivier Ledoit and Michael Wolf
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###########################################################################

using LinearAlgebra
using Statistics

# comments marked with a `#` come from their code
# comments marked with a `#<#>` are ours

function matlab_ledoitwolf_covcor(z)
    x = copy(z)

    # de-mean returns
    t, n  = size(x)
    meanx = mean(x, dims=1)
    x -= repeat(meanx, outer=(t, 1))

    # compute sample covariance matrix
    sample = (1/t) * (x'*x)

    # compute prior
    diagsample = diag(sample)
    sqrtvar = sqrt.(diagsample)

    sqrtvar_mat = repeat(sqrtvar, outer=(1, n))

    rBar = (sum(sample./(sqrtvar_mat .* sqrtvar_mat'))-n)/(n*(n-1))
    prior = rBar * sqrtvar_mat .* sqrtvar_mat'

    #<#> replace the diagonal
    prior -= Diagonal(diag(prior))
    prior += Diagonal(diagsample)

    #<#> compute optimal shrinkage

    # what we call pi-hat
    y = x.^2;
    phiMat = y'*y/t - 2*(x'*x) .* sample/t + sample.^2
    phi = sum(phiMat)

    # what we call rho-hat
    term1    = ((x.^3)'*x)/t
    help     = x'*x/t
    helpDiag = diag(help)
    term2    = repeat(helpDiag, outer=(1, n)) .* sample
    term3    = help .* repeat(diagsample, outer=(1, n))
    term4    = repeat(diagsample, outer=(1, n)) .* sample
    thetaMat = term1 - term2 - term3 + term4
    #<#> remove the diagonal
    thetaMat -= Diagonal(diag(thetaMat))
    rho       = sum(diag(phiMat)) + rBar*sum(sqrtvar'./sqrtvar.*thetaMat)

    # what we call gamma hat
    gamma = sum((sample - prior).^2)
    kappa = (phi - rho) / gamma
    shrinkage = max(0, min(1, kappa/t))

    part_results = Dict(
        "r̄" => rBar,
        "F" => prior,
        "shrinkage" => shrinkage,
        "kappa" => kappa,
        "gamma" => gamma,
        "lwcov" => shrinkage*prior + (1-shrinkage)*sample)

    return part_results
end


export Cubic, Spline, npatches, funk

# see http://mathworld.wolfram.com/CubicSpline.html

import Base: getindex, show

"""
`Cubic(a,b,c,d)` is a 3rd degree polynomial `a+bx+cx^2+dx^3`.
Use `f(x)` to evalue `f` at the value `x`.
"""
struct Cubic
    a::Number
    b::Number
    c::Number
    d::Number
end

(f::Cubic)(x::Number) = f.a + x*(f.b + x*(f.c+x*f.d))

"""
`f'` where `f` is a `Cubic` or `Spline` is the derivative
of `f`.
"""
Base.adjoint(f::Cubic) = Cubic(f.b,2f.c,3f.d,0)

struct Spline
    patches::Array{Cubic,1}
    closed::Bool
end

is_closed(S::Spline) = S.closed

function Base.adjoint(S::Spline)
    plist = adjoint.(S.patches)
    return Spline(plist,S.closed)
end

"""
`Spline(vals,kind)` returns a cubic spline based on the values in `vals`.
The resulting spline `S` will have the property that `S(1)==y[1]`,
`S(2)==y[2]`, and so on up to `S(n)==y[n]` where `n` is the length of `y`.

+ If `kind` is `:open` then the second derivatives at the end points will be zero. (This is the default.)
+ If `kind` is `:closed` then we assume that we are interpolating a periodic function where `S(n+1)==S(1)`.
"""
function Spline(y::Array{T,1},kind::Symbol=:open)::Spline where T<:Number
    n = length(y)
    if kind == :open
        @assert n>2 "Open splines must have at least three points"
        return open_spline(y)
    end
    if kind == :closed
        @assert n>3 "Closed splines must have at least four points"
        return closed_spline(y)
    end

    error("Spline type must be :open or :closed, not $kind")
end

function npatches(S::Spline)
    return length(S.patches)
end

function show(io::IO, S::Spline)
    adjective = is_closed(S) ? "Closed" : "Open"
    print(io, "$adjective spline with $(npatches(S)) patches")
end

getindex(S::Spline, idx::Int) = S.patches[idx]

"""
`funk(S)` converts the Spline `S` into a callable function
(e.g., that can be passed to `plot`).
"""
function funk(S)::Function
    return x -> S(x)
end

function (S::Spline)(x::Real)
    np = npatches(S)

    if is_closed(S)

        p = Int(floor(x))
        x = mod(x,np)
        p = Int(floor(x))

        if p==0
            p = np
            x += np
        end
        f = S[p]
        return f(x-p)
    end

    # open spline
    p = Int(floor(x))
    if x < 1
        p = 1
    end

    if p > np
        p = np
    end
    f = S[p]
    return f(x-p)
end

function open_spline(y::Array{T,1})::Spline where T<:Number
    n = length(y)
    M = zeros(n,n)
    for i=1:n
        M[i,i] = 4
    end
    M[1,1] = 2
    M[n,n] = 2

    for i=1:n-1
        M[i,i+1] = 1
        M[i+1,i] = 1
    end

    rhs = zeros(T,n)
    rhs[1] = 3*(y[2]-y[1])
    for k=2:n-1
        rhs[k] = 3*(y[k+1]-y[k-1])
    end
    rhs[n] = 3*(y[n]-y[n-1])

    D = M\rhs

    a = zeros(Number,n-1)
    b = zeros(Number,n-1)
    c = zeros(Number,n-1)
    d = zeros(Number,n-1)

    for j=1:n-1
        a[j] = y[j]
        b[j] = D[j]
        c[j] = 3*(y[j+1]-y[j])-2D[j]-D[j+1]
        d[j] = 2*(y[j]-y[j+1])+D[j]+D[j+1]
    end
    return Spline( [ Cubic(a[i],b[i],c[i],d[i]) for i=1:n-1 ] , false)
end

function closed_spline(y::Array{T,1})::Spline where T <: Number
    yy = copy(y)
    n = length(y)
    prepend!(yy,y[end-3:end])
    append!(yy,y[1:4])

    S = open_spline(yy)

    plist = S.patches[5:end-3]
    return Spline(plist,true)

end


using DataStructures

# The implementation is referencing the wikipedia page
# https://en.wikipedia.org/wiki/Matching_pursuit#The_algorithm

const default_max_iter = 20
const default_tolerance = 1e-6

function SparseArrays.sparsevec(d::DefaultDict, m::Int)
    SparseArrays.sparsevec(collect(keys(d)), collect(values(d)), m)
end

function matching_pursuit_(data::AbstractVector, dictionary::AbstractMatrix,
                           max_iter::Int, tolerance::Float64)
    n_atoms = size(dictionary, 2)

    residual = copy(data)

    xdict = DefaultDict{Int, Float64}(0.)
    for i in 1:max_iter
        if norm(residual) < tolerance
            return sparsevec(xdict, n_atoms)
        end

        # find an atom with maximum inner product
        products = dictionary' * residual
        _, maxindex = findmax(abs.(products))
        maxval = products[maxindex]
        atom = dictionary[:, maxindex]

        # c is the length of the projection of data onto atom
        a = maxval / sum(abs2, atom)  # equivalent to maxval / norm(atom)^2
        residual -= atom * a

        xdict[maxindex] += a
    end
    return sparsevec(xdict, n_atoms)
end

"""
    matching_pursuit(data::Vector, dictionary::AbstractMatrix;
                     max_iter::Int = $default_max_iter,
                     tolerance::Float64 = $default_tolerance)

Find ``x`` such that ``Dx = y`` or ``Dx ≈ y`` where y is `data` and D is `dictionary`.
```
# Arguments
* `max_iter`: Hard limit of iterations
* `tolerance`: Exit when the norm of the residual < tolerance
```
"""
function matching_pursuit(data::AbstractVector, dictionary::AbstractMatrix;
                          max_iter::Int = default_max_iter,
                          tolerance = default_tolerance)

    if tolerance <= 0
        throw(ArgumentError("`tolerance` must be > 0"))
    end

    if max_iter <= 0
        throw(ArgumentError("`max_iter` must be > 0"))
    end

    if size(data, 1) != size(dictionary, 1)
        throw(ArgumentError(
            "Dimensions must match: `size(data, 1)` and `size(dictionary, 1)`."
        ))
    end

    matching_pursuit_(data, dictionary, max_iter, tolerance)
end

"""
    matching_pursuit(data::AbstractMatrix, dictionary::AbstractMatrix;
                     max_iter::Int = $default_max_iter,
                     tolerance::Float64 = $default_tolerance)

Find ``X`` such that ``DX = Y`` or ``DX ≈ Y`` where Y is `data` and D is `dictionary`.
```
# Arguments
* `max_iter`: Hard limit of iterations
* `tolerance`: Exit when the norm of the residual < tolerance
```
"""
function matching_pursuit(data::AbstractMatrix, dictionary::AbstractMatrix;
                          max_iter::Int = default_max_iter,
                          tolerance::Float64 = default_tolerance)
    K = size(dictionary, 2)
    N = size(data, 2)

    X = spzeros(K, N)

    for i in 1:N
        X[:, i] = matching_pursuit(
            vec(data[:, i]),
            dictionary,
            max_iter = max_iter,
            tolerance = tolerance
        )
    end
    return X
end


function Newton_Spectral(nlp     :: AbstractNLPModel;
                         x       :: AbstractVector=copy(nlp.meta.x0),
                         ϵ       :: Real=√eps(eltype(x)),
                         #ϵ       :: T = 1e-6,
                         maxiter :: Int = 200,
                         Lp      :: Real = 2 # norm Lp
                         )# where T<:Real

    T = eltype(x)
    n = nlp.meta.nvar
    @info log_header([:iter, :f, :dual, :step, :slope], [Int, T, T, T, T],
                     hdr_override=Dict(:f=>"f(x)", :dual=>"‖∇f‖", :slope=>"∇fᵀd"))
    f  ::T = obj(nlp,x)
    ∇f :: Vector{T} = grad(nlp, x)

    τ₀ = 0.0005
    τ₁ = 0.999

    iter = 0
    @info log_row(Any[iter, f, norm(∇f, Lp)])

    while (norm(∇f, Lp) > ϵ) && (iter <= maxiter)
        H = Matrix(Symmetric(hess(nlp, x),:L))
        Δ, O = eigen(H)
        # Boost negative values of Δ to 1e-8
        # devise an adaptative value for γ
        γ = 1e-6
        D = abs.(Δ) + max.((γ .- abs.(Δ)), 0.0) .*ones(n)

        d = - O*diagm(1.0 ./ D)*O'*∇f

        #------------------------------------------
        # Hard coded line search
        hp0 = ∇f'*d
        t=1.0
        # Simple Wolfe forward tracking
        xt = x + t*d
        ∇ft = grad(nlp,xt)
        hp = ∇ft'*d
        ft = obj(nlp, xt)
        #  while  ~wolfe & armijo
        nbW = 0
        while (hp <= τ₁ * hp0) && (ft <= ( f + τ₀*t*hp0)) && (nbW < 10)
            t *= 5
            xt = x + t*d
            ∇ft = grad(nlp,xt)
            hp = ∇ft'*d
            ft = obj(nlp, xt)
            nbW += 1
            @debug "W", ft
        end
        tw = t

        # Simple Armijo backtracking
        nbk = 0
        while (ft > ( f + τ₀*t*hp0)) && (nbk < 20)
            t *= 0.5
            xt = x + t*d
            ft = obj(nlp, xt)
            nbk += 1
            @debug "A", ft
        end
        #------------------------------------------

        x += t*d
        f = ft
        if t!=tw   ∇ft = grad(nlp, xt) end
        ∇f = ∇ft
        iter += 1

        @info log_row(Any[iter, f, norm(∇f, Lp), t, hp0])
    end
    if iter > maxiter @warn "Maximum d'itérations"
    end

    return iter, f, norm(∇f, Lp), x
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.3
#     language: julia
#     name: julia-1.6
# ---

# %%
using LinearAlgebra
using Random
using Distributions
using StatsPlots

# %%
A = [
    2 -1
    -1 2
]

# %%
dist = MvNormal(A)

# %%
A[1,2]/√(A[1,1]*A[2,2])

# %%
n = 2^10
XY = rand(dist, n)
X, Y = XY[1,:], XY[2,:]
@show cor(X, Y)

scatter(X, Y)

# %%
function sim(A, n, L=10^5)
    dist = MvNormal(A)
    R = zeros(L)
    XY = zeros(2, n)
    for i in eachindex(R)
        rand!(dist, XY)
        R[i] = @views cor(XY[1,:], XY[2,:])
    end
    R
end

# %%
n = 2^6
R = sim(A, n)
Z = atanh.(R)
r0 = A[1,2]/√(A[1,1]*A[2,2])
z0 = atanh(r0)
histogram(Z; norm=true, alpha=0.3, bin=200)
plot!(Normal(z0, 1/√(n - 3)))
plot!(fit_mle(Normal, Z))

# %%
n = 2^8
x, y = rand(Uniform(-1, 1), n), rand(Uniform(-1, 1), n)
X, Y = x, x + y
@show cor(X, Y)
scatter(X, Y; label="")

# %%
function sim2(n, L=10^5)
    R = zeros(L)
    d = Uniform(-1, 1)
    x = zeros(n)
    y = similar(x)
    Y = similar(x)
    for i in eachindex(R)
        rand!(d, x)
        rand!(d, y)
        @. Y = x + y
        R[i] = cor(x, Y)
    end
    R
end

# %%
n = 2^6
R = sim2(n)
Z = atanh.(R)
r0 = 1/√2
z0 = atanh(r0)
histogram(Z; norm=true, alpha=0.3, bin=200)
plot!(Normal(z0, 1/√(n - 3)))
plot!(fit_mle(Normal, Z))

# %%
n = 2^8
x, y = rand(Exponential(), n), rand(Exponential(), n)
X, Y = x, x + y
@show cor(X, Y)
scatter(X, Y; label="")

# %%
function sim3(n, L=10^5)
    R = zeros(L)
    d = Exponential()
    x = zeros(n)
    y = similar(x)
    Y = similar(x)
    for i in eachindex(R)
        rand!(d, x)
        rand!(d, y)
        @. Y = x + y
        R[i] = cor(x, Y)
    end
    R
end

# %%
n = 2^6
R = sim3(n)
Z = atanh.(R)
r0 = 1/√2
z0 = atanh(r0)
histogram(Z; norm=true, alpha=0.3, bin=200)
plot!(Normal(z0, 1/√(n - 3)))
plot!(fit_mle(Normal, Z))

# %%
function simplot(dist, n, L=10^5)
    Z = zeros(L)
    x = zeros(n)
    y = similar(x)
    Y = similar(x)
    for i in eachindex(R)
        rand!(dist, x)
        rand!(dist, y)
        @. Y = x + y
        Z[i] = atanh(cor(x, Y))
    end
    r0 = 1/√2
    z0 = atanh(r0)
    histogram(Z; norm=true, alpha=0.3, bin=200)
    plot!(Normal(z0, 1/√(n - 3)))
    plot!(fit_mle(Normal, Z))
end

# %%
simplot(Normal(), 2^6)

# %%
simplot(Uniform(), 2^6)

# %%
simplot(Exponential(), 2^6)

# %%
simplot(MixtureModel([Normal(), Normal(10, 1)], [0.95, 0.05]), 2^4)

# %%
simplot(MixtureModel([Normal(), Normal(10, 1)], [0.95, 0.05]), 2^6)

# %%


# Some **really cool** maths:
#
# $$ \exp(i\pi) + 1 \quad = \quad 0 $$
#
# We can show this with some code:

x = exp(im*π) + 1

# that looks close to zero but

x ≈ 0

# however

abs(x) < eps()

# #### Conclusion
#
# The equation is proven thanks to our very rigorous proof.


using SpecialFunctions

# from datetime import date
# User-provided date:
#
stosym = "wmt"
exyear = 2019
exmonth = 5
exday = 16
stopr = 92.68
strike = 94
callBid = 0.97
callAsk = 1.01
rfir = 0.0225
#
#  Initialize key variables below
#  (This is not actually required. This is your Prof's habit).
#
daystd = 1
cipd = 0.0
d1 = 0.0
durvol = 0.0
cumd1 = 0.0
cumd2 = 0.0
newcp = 0.0
tempcp = 0.0
timedecay = 0.0
tdprice = 0.0
callpr = 0.0
#
# Calculating strike value (PEG), which is somewhere between BB and BA.
# PEG coefficient is explained in lecture.
peg_coef = 0.60
spread = callAsk - callBid
callpr = callBid + (peg_coef*spread)
# Calculating days to expiry:
tnow = Dates.today()
expiry = Dates.Date(exyear, exmonth, exday)
days = Int(expiry - tnow)
#  Our method for calculating the cumulative standard normal distribution:
function csnd(dval)
	return (1.0 + erf(dval/sqrt(2.0)))/2.0
end
#
#  Below we calculate implied daily volatility using a recursive method to
#  converge to an answer.  First we make an initializing calculation, then
#  we converge inside the while loop. The convergence method shown is Prof E's
#  "Divide and Conquer", which is relatively efficient, usually using only 8
#  (+/- 2) steps. The student should recognize that the cumd1 is the delta.
#  The implied daily volatility that we are seeking is "cipd" below.
#
target = callpr
precision = 1e-4
count = 0
low = 0.0
high = 1.0
cipd = (high+low)/2
d1 = log(stopr/strike)+((rfir/365)+(cipd^2)/2)*days
durvol = cipd*sqrt(days)
cumd1 = csnd(d1/durvol)
cumd2 = csnd((d1/durvol) - durvol)
discount = exp(-rfir*days/365)
tempcp = (stopr*cumd1)-(strike*discount*cumd2)
while tempcp<=(target-precision) || tempcp>=(target+precision)
	if tempcp >= (target+precision)
		high = cipd
	else
		low = cipd
    end
	cipd = (high+low)/2
	d1 = log(stopr/strike)+((rfir/365)+(cipd^2)/2)*days
	durvol = cipd*sqrt(days)
	cumd1 = csnd(d1/durvol)
	cumd2 = csnd((d1/durvol) - durvol)
	discount = exp(-rfir*days/365)
	tempcp = (stopr*cumd1)-(strike*discount*cumd2)
	count +=1
end
#	Below we calculate one day time decay using our new value for volatility
days = days - daystd
d1 = log(stopr/strike)+((rfir/365)+(cipd^2)/2)*days
durvol = cipd*sqrt(days)
cumd1 = csnd(d1/durvol)
cumd2 = csnd((d1/durvol) - durvol)
discount = exp(-rfir*days/365)
newcp = (stopr*cumd1)-(strike*discount*cumd2)
timedecay = callpr - newcp

println( "")
println( "Days to expiry: ", days)
println( "Stock price: ", stopr)
println( "Strike price: ", strike)
println( "Call ASK: ", callAsk)
println( "Call BID: ", callBid)
println( "Call price (PEG): ", callpr)
println( "The Delta:", cumd1)
println( "One day time decay:", timedecay)
println( "The call's implied volatility: ", cipd)


#=
In England the currency is made up of pound, £, and pence, p, and there are eight coins in general circulation:

1p, 2p, 5p, 10p, 20p, 50p, £1 (100p) and £2 (200p).
It is possible to make £2 in the following way:

1×£1 + 1×50p + 2×20p + 1×5p + 1×2p + 3×1p
How many different ways can £2 be made using any number of coins?
=#
coins = [1,2,5,10,20,50,100,200]
# Number of ways of making change for n cents
# using only the first k types of coins
function coin(n,k)
  if k < 1 || n < 0
    return 0
  elseif n == 0
    return 1
  end
  coin(n, k-1)+coin(n-coins[k], k)
end
@time println(coin(200,8))


using OrdinaryDiffEq, Plots
gr()

#Half-life of Carbon-14 is 5,730 years.
C₁ = 5.730

#Setup
u₀ = 1.0
tspan = (0.0, 1.0)

#Define the problem
radioactivedecay(u,p,t) = -C₁*u

#Pass to solver
prob = ODEProblem(radioactivedecay,u₀,tspan)
sol = solve(prob,Tsit5())

#Plot
plot(sol,linewidth=2,title ="Carbon-14 half-life", xaxis = "Time in thousands of years", yaxis = "Percentage left", label = "Numerical Solution")
plot!(sol.t, t->exp(-C₁*t),lw=3,ls=:dash,label="Analytical Solution")


#Create Initially traveling gaussian
function get_gaussian(p::Param , Opp::Opps)

    #Get Variables
    Nx  = p.xnodes
    Ny  = p.ynodes

    A0  = p.A0
    σ   = p.σ
    r0  = p.r0
    ω   = p.ω
    m   = p.m

    #Get operators
    x   = Opp.xcoord
    y   = Opp.ycoord

    Dx  = Opp.Dx

    #Create initial state
    ψ_matrix = zeros( Nx , Ny )
    dψ_matrix = zeros( Nx , Ny )

    for i in 1:Nx , j in 1:Ny

        #Radial and angular coordinates
        r  = sqrt(x[i]^2 + y[j]^2)
        θ   = atan(y[j] , x[i])

        ψ_matrix[i,j] = A0*cos( m*θ ) * sin( ω*r ) * exp( - 0.5* ( (r-r0)/ σ)^2 )
        dψ_matrix[i,j] = A0*cos( m*θ ) * exp( - 0.5* ( (r-r0)/ σ)^2 ) * ( ω*cos(ω*r) - (r-r0)*sin(ω*r)/σ^2  )

    end

    ψ  = reshape(ψ_matrix , Nx*Ny)
    dψ = reshape(dψ_matrix , Nx*Ny)

    #Return initial state
    return ψ , dψ

end


"""
    cartesian_to_lat_lon(x, y, z)

Convert 3D coordinates (x, y, z) from the unit sphere to (lat, lon). Assumes "lat" is positive
with "z", equatorial plane falls at z=0  and "lon" is measured anti-clockwise (eastward)
from x-axis (y=0) about z-axis.

This is a Julia translation of MATLAB code from MITgcm [1].

[1]: http://wwwcvs.mitgcm.org/viewvc/MITgcm/MITgcm_contrib/high_res_cube/matlab-grid-generator/map_xyz2lonlat.m?view=markup
"""
cartesian_to_lat_lon(x, y, z) = cartesian_to_latitude(x, y, z), cartesian_to_longitude(x, y, z)

cartesian_to_latitude(x, y, z) = atand(z, hypot(x, y))

cartesian_to_longitude(x, y, z) = atand(y, x)


#=
The sum of the squares of the first ten natural numbers is,

12 + 22 + ... + 102 = 385
The square of the sum of the first ten natural numbers is,

(1 + 2 + ... + 10)2 = 552 = 3025
Hence the difference between the sum of the squares of the first ten natural numbers and the square of the sum is 3025 − 385 = 2640.

Find the difference between the sum of the squares of the first one hundred natural numbers and the square of the sum.
=#

function sumsquaredifferences()
  collection = 1:100
  x = sum(map((x)->x^2, collection))
  y = sum(collection)^2
  y - x
end
@time println(sumsquaredifferences())


using EconPDEs, Distributions

# Arbitrage With Holding Costs: A Utility-Based Approach (Bruce Tuckman and Jean-Luc Vila)
Base.@kwdef struct ArbitrageHoldingCosts
    c::Float64 = 0.06
    r::Float64 = 0.09
    ρ::Float64 = 5.42
    σ::Float64 = 26.72
    a::Float64 = 26.72
    T::Float64 = 100
end

function (m::ArbitrageHoldingCosts)(state::NamedTuple, y::NamedTuple, τ::Number)
    (; c, r, ρ, σ, a, T) = m
    (; z) = state
    (; F, Fz_up, Fz_down, Fzz) = y
    Fz = (z >= 0) ? Fz_up : Fz_down
    ϕ = z * (1 + z^2)^(-1/2)
    ϕz = (1 + z^2)^(-3/2)
    ϕzz = - 3 * z * (1 + z^2)^(-5/2)
    sτ = c / r * (1 - exp(- r * (T - τ)))
    μL = ρ * z - 0.5 * σ^2 * ϕzz / ϕz + c / sτ * (ϕ - 1) / ϕz
    iL = 1 / a * (μL / σ^2 + Fz)
    μS = ρ * z - 0.5 * σ^2 * ϕzz / ϕz + c / sτ * (ϕ + 1) / ϕz
    iS = 1 / a * (μS / σ^2 + Fz)
    μ, i  = 0.0, 0.0
    if iL > 0
        i = iL
        μ = μL
    elseif iS < 0
        i = iS
        μ = μS
    end
    Ft = - ((μ + σ^2 * Fz) * a * i - 0.5 * σ^2 * (a * i)^2 - ρ * z * Fz + 0.5 * σ^2 * (Fzz - Fz^2))
    return (; Ft)
end

m = ArbitrageHoldingCosts()
zn = 200
d = Normal(0, sqrt(m.σ^2 / (2 * m.ρ)))
stategrid = OrderedDict(:z => range(quantile(d, 0.00001), quantile(d, 0.99999), length = zn))
yend = OrderedDict(:F => zeros(zn))
τs = range(0, m.T, length = 100)
result = pdesolve(m, stategrid, yend, τs)
residual_norm = maximum(result.residual_norm)

### reproduce Fig 2
#d = Normal(0, sqrt(m.σ^2 / (2 * m.ρ)))
#zmin = quantile(d, 0.025)
#zmax = quantile(d, 0.975)
#idx = (state[:z] .>= zmin) .& (state[:z] .<= zmax)
#
#using Plots
#plot(result[:x][idx, 20], [result[:I_myopic][idx, 2] result[:I][idx, 2]], label = ["myopic" "all"])
#plot(result[:x][idx, 20], result[:I][idx, 2])
#plot!(result[:x][idx, 60], result[:I][idx, 60])
#plot!(result[:x][idx, 100], result[:I][idx, 100])
#


# A quadratic segment in 3D space that passes through three points: x⃗₁, x⃗₂, and x⃗₃.
# The assumed relation of the points may be seen in the diagram below.
#                 ___x⃗₃___
#            ____/        \____
#        ___/                  \
#     __/                       x⃗₂
#   _/
#  /
# x⃗₁
#
# NOTE: x⃗₃ is between x⃗₁ and x⃗₂, but not necessarily the midpoint.
# q(r) = (2r-1)(r-1)x⃗₁ + r(2r-1)x⃗₂ + 4r(1-r)x⃗₃
# See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
# Chapter 8, Advanced Data Representation, in the interpolation functions section
struct QuadraticSegment_3D{T <: AbstractFloat}
    points::NTuple{3, Point_3D{T}}
end

# Constructors
# -------------------------------------------------------------------------------------------------
QuadraticSegment_3D(p₁::Point_3D{T},
                    p₂::Point_3D{T},
                    p₃::Point_3D{T}) where {T <: AbstractFloat} = QuadraticSegment_3D((p₁, p₂, p₃))

# Methods
# -------------------------------------------------------------------------------------------------
function (q::QuadraticSegment_3D{T})(r::R) where {T <: AbstractFloat, R <: Real}
    # See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
    # Chapter 8, Advanced Data Representation, in the interpolation functions section
    rₜ = T(r)
    return (2rₜ-1)*(rₜ-1)*q.points[1] + rₜ*(2rₜ-1)*q.points[2] + 4rₜ*(1-rₜ)*q.points[3]
end

#function derivative(q::QuadraticSegment_3D{T}, r::R) where {T <: AbstractFloat, R <: Real}
#    # dq⃗/dr
#    rₜ = T(r)
#    return (4rₜ - 3)*q.points[1] + (4rₜ - 1)*q.points[2] + (4 - 8rₜ)*q.points[3]
#end
#
#function arc_length(q::QuadraticSegment_3D{T}; N::Int64=20) where {T <: AbstractFloat}
#    # Mathematica solution is pages long and can produce NaN results when the segment is
#    # straight, so numerical integration is used. (Gauss-Legengre quadrature)
#    #     1                  N
#    # L = ∫ ||q⃗'(r)||dr  ≈   ∑ wᵢ||q⃗'(rᵢ)||
#    #     0                 i=1
#    # The default number of points is N = 20, since the timing difference is very small for
#    # additional accuracy when compared with N = 15, and small N give poor accuracy.
#    w, r = gauss_legendre_quadrature(T, N)
#    return sum(norm.(w .* derivative.(q, r)))
#end
#
#function intersect(l::LineSegment_3D{T}, q::QuadraticSegment_3D{T}) where {T <: AbstractFloat}
#    # q(r) = (2r-1)(r-1)x⃗₁ + r(2r-1)x⃗₂ + 4r(1-r)x⃗₃
#    # q(r) = 2r²(x⃗₁ + x⃗₂ - 2x⃗₃) + r(-3x⃗₁ - x⃗₂ + 4x⃗₃) + x⃗₁
#    # Let D⃗ = 2(x⃗₁ + x⃗₂ - 2x⃗₃), E⃗ = (-3x⃗₁ - x⃗₂ + 4x⃗₃), F⃗ = x₁
#    # q(r) = r²D⃗ + rE⃗ + F⃗
#    # l(s) = x⃗₄ + sw⃗
#    # If D⃗ × w⃗ ≠ 0⃗
#    #   x⃗₄ + sw⃗ = r²D⃗ + rE⃗ + F⃗
#    #   sw⃗ = r²D⃗ + rE⃗ + (F⃗ - x⃗₄)
#    #   0⃗ = r²(D⃗ × w⃗) + r(E⃗ × w⃗) + (F⃗ - x⃗₄) × w⃗
#    #   Let A⃗ = (D⃗ × w⃗), B⃗ = (E⃗ × w⃗), C⃗ = (F⃗ - x⃗₄) × w⃗
#    #   0⃗ = r²A⃗ + rB⃗ + C⃗
#    #   0 = (A⃗ ⋅ A⃗)r² + (B⃗ ⋅ A⃗)r + (C⃗ ⋅ A⃗)
#    #   A = (A⃗ ⋅ A⃗), B = (B⃗ ⋅ A⃗), C = (C⃗ ⋅ A⃗)
#    #   0 = Ar² + Br + C
#    #   r = (-B - √(B²-4AC))/2A, -B + √(B²-4AC))/2A)
#    #   s = ((q(r) - p₄)⋅w⃗/(w⃗ ⋅ w⃗)
#    #   r is invalid if:
#    #     1) A = 0
#    #     2) B² < 4AC
#    #     3) r < 0 or 1 < r   (Curve intersects, segment doesn't)
#    #   s is invalid if:
#    #     1) s < 0 or 1 < s   (Line intersects, segment doesn't)
#    # If D⃗ × w⃗ = 0, we need to use line intersection instead.
#    bool = false
#    npoints = 0
#    points = [Point_3D(T, 0), Point_3D(T, 0)]
#    D⃗ = 2*(q.points[1] + q.points[2] - 2*q.points[3])
#    E⃗ = 4*q.points[3] - 3*q.points[1] - q.points[2]
#    w⃗ = l.points[2] - l.points[1]
#    A⃗ = D⃗ × w⃗
#    B⃗ = E⃗ × w⃗
#    C⃗ = (q.points[1] - l.points[1]) × w⃗
#    A = A⃗ ⋅ A⃗
#    B = B⃗ ⋅ A⃗
#    C = C⃗ ⋅ A⃗
#    if abs(A) < 1.0e-6
#        # Line intersection
#        r = (-C⃗ ⋅ B⃗)/(B⃗ ⋅ B⃗)
#        s = (q(r)- l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#        points[1] = q(r)
#        if (0 ≤ s ≤ 1) && (0 ≤ r ≤ 1)
#            bool = true
#            npoints = 1
#        end
#    elseif B^2 ≥ 4A*C
#        # Quadratic intersection
#        r₁ = (-B - √(B^2-4A*C))/2A
#        r₂ = (-B + √(B^2-4A*C))/2A
#        points[1] = q(r₁)
#        points[2] = q(r₂)
#        s₁ = (points[1] - l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#        s₂ = (points[2] - l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#
#        # Check points to see if they are valid intersections.
#        # First r,s valid?
#        if (0 ≤ r₁ ≤ 1) && (0 ≤ s₁ ≤ 1) && (points[1] ≈ l(s₁))
#            npoints += 1
#        end
#        # Second r,s valid?
#        if (0 ≤ r₂ ≤ 1) && (0 ≤ s₂ ≤ 1) && (points[2] ≈ l(s₂))
#            npoints += 1
#            # If only point 2 is valid, return it in index 1 of points
#            if npoints === 1
#                points[1] = points[2]
#            end
#        end
#        bool = npoints > 0
#    end
#    return bool, npoints, points
#end
#intersect(q::QuadraticSegment_3D, l::LineSegment_3D) = intersect(l, q)


function combined(op, kernel_function)
    function combined_kernel_function(u, v, k=nothing)
        ker = kernel_function(u, v)
        isnothing(k) ? ker : op.(k, ker)
    end
end

# Kernel functions

"""
    gaussiankernel(u, v)

Compute the Gaussian kernel, given by the formula
```
    κ(x,y) = exp(-‖x-y‖²/2).
```
`u` and `v` are matrices, where each column is a datapoint.
"""
function gaussiankernel(u, v)
    uu = sum(abs2, u, dims=1)
    vv = sum(abs2, v, dims=1)
    acc = u' * v
    @. acc = exp(acc - uu' / 2 - vv / 2)
    return acc
end

function rrule(::typeof(gaussiankernel), u, v)
    r = gaussiankernel(u, v)
    # r̄ is short for ∂l / ∂r
    function gaussiankernel_pullback(r̄)
        m = r .* r̄ # pullback the exponential
        ū = v * m'
        ū .-= sum(m', dims=1) .* u
        v̄ = u * m
        v̄ .-= sum(m, dims=1) .* v
        return NO_FIELDS, ū, v̄
    end
    return r, gaussiankernel_pullback
end

"""
    additivegaussiankernel(u, v, k=nothing)

Compute the radial basis function kernel (see [`gaussiankernel`](@ref)) on `u`, `v`.
If `k` is not `nothing`, add it to the result.
"""
const additivegaussiankernel = combined(+, gaussiankernel)

"""
    additivegaussiankernel(u, v, k=nothing)

Compute the radial basis function kernel (see [`gaussiankernel`](@ref)) on `u`, `v`.
If `k` is not `nothing`, multiply the result by it.
"""
const multiplicativegaussiankernel = combined(*, gaussiankernel)


# See https://en.wikipedia.org/wiki/Fibonacci

fibonacci(n) = ([1 1 ; 1 0] ^ n)[1, 2]

expressions = [
    :(fibonacci(0))
    :(fibonacci(1))
    :(fibonacci(6))
    :(fibonacci(91))
    :(fibonacci(big(283)))
    :(fibonacci(big(3674)))
]

for expr in expressions
    println("[] :julia-answer \"\"\"", escape_string(string(expr)), " = ", eval(expr), "\"\"\".")
end


export lambertW

#
# The Lambert W function (See reference [4])
# The following algorithm computes the Lambert-W function to machine accuracy using Halley's method for any x
# in the domain I=( -1/exp(1), ∞ ). This function can also compute the lambert-W function component-wise for
# vectors or matrices. The Lambert-W function is used in the module DESincEig to compute the mesh size h.
# Input: x:: Number, Vector or Matrix
# Output: lambertW(x)
#
function lambertW(x::Real)
    if x < -exp(-one(x))
        return throw(DomainError())
    elseif x == -exp(-one(x))
        return -one(x)
    elseif x < 0
        w0 = e*x/(1+inv(inv(sqrt(2*e*x+2))+inv(e-1)-inv(sqrt(2))))
    else
        logx = log(1+x)
        w0 = logx*(1-log(1+logx)/(2+logx))
    end
    expw0 = exp(w0)
    w1 = w0 - (w0*expw0 - x)/((w0 + 1)*expw0 -
        (w0 + 2) * (w0*expw0 - x)/(2w0 + 2))
    while abs(w1/w0 - 1) > 2eps(typeof(x))
        w0 = w1
        expw0 = exp(w0)
        w1 = w0 - (w0*expw0 - x)/((w0 + 1)*expw0 -
            (w0 + 2) * (w0*expw0 - x)/(2w0 + 2))
    end
    return w1
end
lambertW(x::Integer) = lambertW(float(x))
@vectorize_1arg Real lambertW


#
# this file shows an example of pseudo-spectral discretization for PHS systems for wave equation
#  three cases are presented. It should be used for verifying the discretization, eigenvalues,
#  and constraint elimination methods. Three cases are tested:
#         1) PHS with u = 0 (Neumann-Dirichlet B.C)
#         2) PHS with y = 0 (Neumann-Dirichlet B.C, but with a constraint)
#         3) PHS with y = 0 after removing the constraints.
#

# auxiliary function: used for printing results

function print_table(table)
	num_y = size(table,1)
	num_x = size(table,2)
	for i = 1:num_y
		for j = 1:num_x
			if j < 3
				@printf("%.2f\t", table[i,j])
			else
				@printf("%.2e", table[i,j])
			end
		end
		@printf("\n")
	end
end

# modules:
using PortHamiltonian
#using PyPlot

N = 10;
a = 0.;
b = 1.;

# load the phs discretization structure
## the Hamiltonian is considered linear with Q = I
ph = discrete_phs(N, a, b);

eigval, eigvec = eig(ph);
num_freq = frequencies(eigval)

exact_freq = pi/2*collect(1:2:2*N);
err = exact_freq - num_freq;

# comparison between numerical and exact results:
comp_table = [num_freq exact_freq err./exact_freq];
comp_table = comp_table[1:10,:]
@printf("\nNatural frequencies of the Dirichlet-Neumann B.C. wave equation (with u = 0)\n")
@printf("exact\tnumeric\terror \n")
print_table(comp_table)
@printf("\t From a total %d frequencies, %d are represented with error less than 0.01\n", length(err), sum(abs(err) .< 0.01))
@printf("\n\t %d are smaller than %.1e\n",sum(abs(err) .< eps()*1e6),eps()*1e6)
#plot the first five modes
#plot(ph.disc_data.flow.xi, real(eigvec[1:N,N+2]))

# constrained version
ph_constrained = deepcopy(ph);
set_constraint!(ph_constrained, ph_constrained.B[:,1], ph_constrained.D[1:1,1:1]);

## eigenvalues of the constrained
eigval, eigvec = eig(ph_constrained)
num_freq = frequencies(eigval)
exact_freq = pi*collect(-0:1:N-1);
err = exact_freq - num_freq;

# comparison between numerical and exact results:
comp_table = [num_freq exact_freq err./exact_freq]
comp_table = comp_table[1:10,:]
@printf("\nNatural frequencies of the Neumann-Dirichlet B.C. wave equation (with y = 0)\n")
@printf("exact\tnumeric\terror \n")
print_table(comp_table)
@printf("\t\t From a total %d frequencies, %d are represented with error less than 0.01\n", length(err), sum(abs(err) .< 0.01))
@printf("\n\t %d are smaller than %.1e\n",sum(abs(err) .< eps()*1e6),eps()*1e6)
## find an explicit equivalent system
ph_c_elim = constraint_elimination(ph_constrained)
eigval, eigvec = eig(ph_c_elim);
num_freq = frequencies(eigval)

exact_freq = pi*collect(0:1:N-1);
err = exact_freq - num_freq;
# comparison between numerical and exact results:
comp_table = [num_freq exact_freq err./exact_freq]
comp_table = comp_table[1:10,:]
@printf("\nNatural frequencies of the Neumann-Dirichlet B.C. wave equation (with y=0 after constraint elimination)\n")
@printf("exact\tnumeric\terror \n")
print_table(comp_table)
@printf("\t\t From a total %d frequencies, %d are represented with error less than 0.01", length(err), sum(abs(err) .< 0.01))
@printf("\n\t %d are smaller than %.1e\n",sum(abs(err) .< eps()*1e6),eps()*1e6)


"""
    InterpolationLineSearch{C, T}

`InterpolationLineSearch` uses quadratic and cubic interpolations to compute the step length
satisfying the Armijo sufficient decrease condition.

# Fields
- `c::C`: the constant for the Armijo condition in the open interval (0, 1)
- `ϵ::T`: the minimum step length decrease from the initial step length
"""
struct InterpolationLineSearch{C, T}
    c::C
    ϵ::T

    function InterpolationLineSearch(c, ϵ)
        0 < c < 1 || error("`c` must be in the open interval (0, 1)")
        return new{typeof(c), typeof(ϵ)}(c, ϵ)
    end
end

"""
    InterpolationLineSearch(; kwargs...)

Initiate `InterpolationLineSearch`.

# Keywords
- `c=1e-4`: the constant for the Armijo condition in the open interval (0, 1)
- `ϵ=1e-6`: the minimum step length decrease from the initial step length
"""
InterpolationLineSearch(; c=1e-4, epsilon=1e-6) = InterpolationLineSearch(c, epsilon)

function (ils::InterpolationLineSearch)(f, state, p, α₀::T) where {T}
    x = state.x
    ϕ₀ = state.f
    ∇y = state.∇f
    dϕ₀ = ∇y⋅p
    αₖ = α₀
    c = ils.c
    ϵ = ils.ϵ
    ϕ(α) = f(x + α * p)
    ϕαₖ = ϕ(αₖ)
    αₖ₊₁ = T(-(dϕ₀ * αₖ^2) / (2 * (ϕαₖ - ϕ₀ - dϕ₀ * αₖ)))
    while f(x + αₖ₊₁*p) > ϕ₀ + c*αₖ₊₁*dϕ₀
        ϕαₖ₊₁ = ϕ(αₖ₊₁)
        a, b = [αₖ^2 -αₖ₊₁^2; -αₖ^3 αₖ₊₁^3] *
            [ϕαₖ₊₁ - ϕ₀ - dϕ₀*αₖ₊₁; ϕαₖ - ϕ₀ - dϕ₀*αₖ] /
            (αₖ^2*αₖ₊₁^2*(αₖ₊₁-αₖ))
        αₖ₊₂ = (-b + sqrt(b^2 - 3a*dϕ₀)) / (3a)
        αₖ₊₂ = max(min(αₖ₊₂, αₖ₊₁ - ϵ), αₖ₊₁/2)
        αₖ = αₖ₊₁
        ϕαₖ = ϕαₖ₊₁
        αₖ₊₁ = T(αₖ₊₂)
    end
    return αₖ₊₁
end


using ReachabilityAnalysis, SparseArrays, Symbolics

const var = @variables x[1:9] t

function platoon_connected(; deterministic_switching::Bool=true,
                             c1=5.0)  # clock constraints
    n = 9 + 1

    # x' = Ax + Bu + c
    A = Matrix{Float64}(undef, n, n)
    A[1, :] = [0, 1.0, 0, 0, 0, 0, 0, 0, 0, 0]
    A[2, :] = [0, 0, -1.0, 0, 0, 0, 0, 0, 0, 0]
    A[3, :] = [1.6050, 4.8680, -3.5754, -0.8198, 0.4270, -0.0450, -0.1942,  0.3626, -0.0946, 0.]
    A[4, :] = [0, 0, 0, 0, 1.0, 0, 0, 0, 0, 0,]
    A[5, :] = [0, 0, 1.0, 0, 0, -1.0, 0, 0, 0, 0]
    A[6, :] = [0.8718, 3.8140, -0.0754,  1.1936, 3.6258, -3.2396, -0.5950,  0.1294, -0.0796, 0.]
    A[7, :] = [0, 0, 0, 0, 0, 0, 0, 1.0, 0, 0]
    A[8, :] = [0, 0, 0, 0, 0, 1.0, 0, 0, -1.0, 0]
    A[9, :] = [0.7132, 3.5730, -0.0964,  0.8472, 3.2568, -0.0876,  1.2726,  3.0720, -3.1356, 0.]
    A[10, :] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0]; # t' = 1

    if deterministic_switching
        invariant = HalfSpace(t <= c1, var)
    else
        invariant = Universe(n)
    end

    # acceleration of the lead vehicle + time
    B = sparse([2], [1], [1.0], n, 1)
    U = Hyperrectangle(low=[-9.], high=[1.])
    c = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1.0]
    @system(x' = Ax + Bu + c, x ∈ invariant, u ∈ U)
end

function platoon_disconnected(; deterministic_switching::Bool=true,
                                c2=5.0)  # clock constraints
    n = 10 # 9 dimensions + time

    # x' = Ax + Bu + c
    A = Matrix{Float64}(undef, n, n)
    A[1, :] = [0, 1.0, 0, 0, 0, 0, 0, 0, 0, 0]
    A[2, :] = [0, 0, -1.0, 0, 0, 0, 0, 0, 0, 0]
    A[3, :] = [1.6050, 4.8680, -3.5754, 0, 0, 0, 0, 0, 0, 0]
    A[4, :] = [0, 0, 0, 0, 1.0, 0, 0, 0, 0, 0,]
    A[5, :] = [0, 0, 1.0, 0, 0, -1.0, 0, 0, 0, 0]
    A[6, :] = [0, 0, 0,  1.1936, 3.6258, -3.2396, 0, 0, 0, 0.]
    A[7, :] = [0, 0, 0, 0, 0, 0, 0, 1.0, 0, 0]
    A[8, :] = [0, 0, 0, 0, 0, 1.0, 0, 0, -1.0, 0]
    A[9, :] = [0.7132, 3.5730, -0.0964,  0.8472, 3.2568, -0.0876,  1.2726,  3.0720, -3.1356, 0.]
    A[10, :] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0]; # t' = 1

    if deterministic_switching
        invariant = HalfSpace(t <= c2, var)
    else
        invariant = Universe(n)
    end

    # acceleration of the lead vehicle + time
    B = sparse([2], [1], [1.0], n, 1)
    U = Hyperrectangle(low=[-9.], high=[1.])
    c = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1.0]
    @system(x' = Ax + Bu + c, x ∈ invariant, u ∈ U)
end

function platoon(; deterministic_switching::Bool=true,
                   c1=5.0,  # clock constraints
                   c2=5.0,  # clock constraints
                   tb=10.0,  # lower bound for loss of communication
                   tc=20.0, tr=20.0) # upper bound for loss of communication (tc) and reset time (tr)

    # three variables for each vehicle, (ei, d(et)/dt, ai) for
    # (spacing error, relative velocity, speed), and the last dimension is time
    n = 9 + 1

    # transition graph
    automaton = LightAutomaton(2)
    add_transition!(automaton, 1, 2, 1)
    add_transition!(automaton, 2, 1, 2)

    # modes
    mode1 = platoon_connected(deterministic_switching=deterministic_switching, c1=c1)
    mode2 = platoon_disconnected(deterministic_switching=deterministic_switching, c2=c2)
    modes = [mode1, mode2]

    # common reset
    reset = Dict(n => 0.)

    # transition l1 -> l2
    if deterministic_switching
        guard = Hyperplane(t == c1, var)
    else
        guard = HPolyhedron([tb <= t, t <= tc], var)
    end
    t1 = ConstrainedResetMap(n, guard, reset)

    # transition l2 -> l1
    if deterministic_switching
        guard = Hyperplane(t == c2, var)
    else
        guard = HalfSpace(t <= tr, var)
    end
    t2 = ConstrainedResetMap(n, guard, reset)
    resetmaps = [t1, t2]

    H = HybridSystem(automaton, modes, resetmaps, [AutonomousSwitching()])

    # initial condition is at the orgin in mode 1
    X0 = BallInf(zeros(n), 0.0)
    initial_condition = [(1, X0)]

    return IVP(H, initial_condition)
end

function dmin_specification(sol, dmin)
    return (-ρ(sparsevec([1], [-1.0], 10), sol) > -dmin) &&
           (-ρ(sparsevec([4], [-1.0], 10), sol) > -dmin) &&
           (-ρ(sparsevec([7], [-1.0], 10), sol) > -dmin)
end

prob_PLAD01 = platoon();

const boxdirs = BoxDirections(10)
const octdirs = OctDirections(10);

length(boxdirs)

alg = BOX(δ=0.01)
sol_PLAD01_BND42 = solve(prob_PLAD01,
                         alg=alg,
                         clustering_method=BoxClustering(1),
                         intersection_method=TemplateHullIntersection(boxdirs),
                         intersect_source_invariant=false,
                         tspan = (0.0 .. 20.0));

dmin_specification(sol_PLAD01_BND42, 42)

-ρ(sparsevec([1], [-1.0], 10), sol_PLAD01_BND42)

-ρ(sparsevec([4], [-1.0], 10), sol_PLAD01_BND42)

-ρ(sparsevec([7], [-1.0], 10), sol_PLAD01_BND42)

using Plots, LaTeXStrings

plot(sol_PLAD01_BND42, vars=(0, 1), xlab=L"t", ylab=L"x_1", title="PLAD01 - BND42", lw=0.1)
plot!(x->x, x->-42., 0., 20., linewidth=2, color="red", linestyle=:dash, leg=nothing)

length(octdirs)

alg = LGG09(δ=0.03, template=octdirs, approx_model=Forward(setops=octdirs));
sol_PLAD01_BND30 = solve(prob_PLAD01,
                         alg=alg,
                         clustering_method=LazyClustering(1),
                         intersection_method=TemplateHullIntersection(octdirs),
                         intersect_source_invariant=false,
                         tspan = (0.0 .. 20.0));

dmin_specification(sol_PLAD01_BND30, 30)

-ρ(sparsevec([1], [-1.0], 10), sol_PLAD01_BND30)

-ρ(sparsevec([4], [-1.0], 10), sol_PLAD01_BND30)

-ρ(sparsevec([7], [-1.0], 10), sol_PLAD01_BND30)

plot(sol_PLAD01_BND30, vars=(0, 1), xlab=L"t", ylab=L"x_1", title="PLAD01 - BND30", lw=0.1)
plot!(x->x, x->-30., 0., 20., linewidth=2, color="red", linestyle=:dash, leg=nothing)


using Plots, LaTeXStrings, Statistics

function Deposition(;len, tot_time, time_steps, rate)
    Time = ceil.(Int, exp.(0:(tot_time)/(time_steps):tot_time))
    surf = [0 for i=1:len]
    VarList = [0.0 for i=1:time_steps]
    for n in 1:time_steps
        randsurf = rand(1:len,(Time[n+1]-Time[n])*rate)
        for i in randsurf
            surf[i] += 1
        end
        VarList[n] = std(surf)
    end
    return VarList
end

function Linear_fit(;len, tot_time, time_steps, rate)
    A = [hcat(log.(Time)) reshape(ones(time_steps), time_steps, 1)]
    b = reshape(log.(meanVar), time_steps, 1)
    line = (A \ b)
    x = 0:tot_time
    y = x .* line[1] .+ line[2]
    return x, y, line
end

iternum = 1000
Parameters = Dict(
                :len => 200,
                :tot_time => 10,
                :time_steps => 20,
                :rate => 10)
allVar = [[0.0 for i in 1:Parameters[:time_steps]] for j = 1:iternum]
meanVar = [0.0 for i in 1:Parameters[:time_steps]]
vars = [0.0 for i in 1:Parameters[:time_steps]]
for i in 1:iternum
    VarList = Deposition(;Parameters...)
    allVar[i] = VarList
    meanVar += VarList
    print("\r$i")
end
meanVar /= iternum
for i in 1:Parameters[:time_steps]
    vars[i] = std(log.(hcat(allVar...))[i,:])
end

Time = ceil.(Int, exp.(0:(Parameters[:tot_time])/(Parameters[:time_steps]):Parameters[:tot_time]))[1:end-1]
X, Y, Line = Linear_fit(;Parameters...)

scatter(log.(Time),log.(meanVar),
    xlabel= L"Log\ Time",
    ylabel= L"Log\ W_{(t)}",
    title= L"Log-Log\ Plot\ of\ ~W_{(t)}-Time~, Particles \ number \approx e^{%$(Parameters[:tot_time])}",
    yerror = vars,
    label = L"Data\ point")
plot!(X,Y,c= :black,label = L"y = %$(round(Line[1],digits= 2))x + %$(round(Line[2],digits= 2))")
savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\ComputationalPhysics-Fall2021\\ProblemSet2\\Figs\\Q2\\RBDW(t).png")


### A Pluto.jl notebook ###
# v0.10.1

using Markdown
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 4a816b10-c188-11ea-3935-1fc7371387ac
using Plots

# ╔═╡ 62a8f092-c198-11ea-2d26-632729a6bd7f
using StatsBase

# ╔═╡ 1c0741f0-c17a-11ea-3c6f-31f0c86dd08d
md"""
# Random Walk Models
"""

# ╔═╡ 58e6d60e-c17c-11ea-145e-ffdc97ab1f3f
md"Short function definition (one-liner):"

# ╔═╡ a93d6520-c17c-11ea-2cc7-5ff51f20df2b
jump() = rand( (-1, +1) )

# ╔═╡ bfa68c10-c17c-11ea-281c-25e48e4ab81d
md"Exactly the same"

# ╔═╡ ccef6d10-c17c-11ea-13e5-432b940ba726
function jump_()
	return rand( (-1, +1) )
end

# ╔═╡ fea8c400-c17c-11ea-35fc-e98ef849c316
jump()

# ╔═╡ 9aff5fa0-c185-11ea-145a-c3230681bd4c
md"It is better to use `tuples` in this case instead of `arrays`, tuples are not mutable and with fixed size"

# ╔═╡ 0762ce10-c17d-11ea-1f2a-172e4ed0500a
begin
	local t = (-1, +1)
	typeof(t)
end

# ╔═╡ 212df5a0-c186-11ea-0f5c-83f0dd5a1497
md"Side remark: Simple bernoulli function"

# ╔═╡ 8ce720ae-c185-11ea-380c-2df4848fb56a
bernoulli(p) = rand() < p

# ╔═╡ f52fc0a0-c185-11ea-2126-b10b11defe13
bernoulli(0.25)

# ╔═╡ 891b2f20-c186-11ea-024f-372f0eef600c
md"## Random walks"

# ╔═╡ 15990630-c186-11ea-2059-e1e301040dae
[jump() for i ∈ 1:10]

# ╔═╡ 62d228a2-c186-11ea-3ced-1db44003e269
function walk(n)
	x = 0

	for i ∈ 1:n
		x += jump() #x = x + jump()
	end

	return x
end

# ╔═╡ 0b66e5a2-c187-11ea-36de-279c4d6aa180
walk(10)

# ╔═╡ 432008a0-c187-11ea-3606-159ae4ed8460
function trajectory(n)
	x = 0

	xs = Int64[x] # inititalize with an initial value and the type of the array

	for i ∈ 1:n
		x += jump() # x = x + jump()

		push!(xs, x)
	end

	return xs
end

# ╔═╡ 131fcc00-c18a-11ea-09dc-0906c430d399
md"""steps = $(@bind steps html"<input type='range' min=5 max=100>" ) """

# ╔═╡ 1afa1280-c18c-11ea-1d1b-7f277d76e68c
md"**$(steps)** steps"

# ╔═╡ 8cfcb9a2-c187-11ea-1129-ff3d55e25de4
trajectories = trajectory(steps)

# ╔═╡ 3b677390-c188-11ea-1079-01f262b1d76f
md"## Plotting"

# ╔═╡ 643c5510-c188-11ea-2d15-fbf4939ae53c
begin
	plot(trajectories, m=:o, size=(300, 200))
	hline!([0], ls=:dash)
end

# ╔═╡ b417eaf0-c18c-11ea-3ae7-19f88eab6b9f
begin
	local num_walkers = 1000
	local num_steps = 100

	local p = plot(size=(500, 400), leg=false)

	for i ∈ 1:num_walkers
		traj = trajectory(num_steps)

		plot!(traj)
	end

	p
end

# ╔═╡ eb35c2e0-c18d-11ea-3209-bd013b156820
md"## Precompute the data"

# ╔═╡ 1ec51ce0-c18f-11ea-2166-6923cbcef091
n = 20

# ╔═╡ 2c106ee0-c18f-11ea-21dc-db7d6706f5cb
traj = trajectory(n)

# ╔═╡ 3ee28c60-c18f-11ea-3297-3d26c8d37627
@bind traj_steps html"<input type='range' min=1 max=20>"

# ╔═╡ 1ac71ed0-c190-11ea-2cfe-4d7cfd514800
md"`traj_steps: ` $(traj_steps)"

# ╔═╡ 727357e0-c18e-11ea-24f4-81a80c9e5f36
begin
	plot(traj[1:traj_steps])
	xlims!(0, length(traj))
	ylims!(minimum(traj), maximum(traj))
end

# ╔═╡ 55e2ff60-c191-11ea-1b6b-b9a74d8fab5f
md"## Run N walkers simultaneously"

# ╔═╡ 60381bd0-c191-11ea-36a6-69333299f23b
md"Rule: NEVER copy and paste (excpet occasionally)

Instead: Automate "

# ╔═╡ e1c193c0-c191-11ea-0c16-67d6d7611c26
begin
	local n = 20

	traj_ = [trajectory(n)]

	push!(traj_, trajectory(n))
end

# ╔═╡ 6fad5d40-c192-11ea-2573-01f0eb6c17ed
md"`append!` extracts the elements of a vector"

# ╔═╡ f4e1b4c0-c192-11ea-3782-79f7c856869f
begin
	local v = [1, 2, 3 ,4]

	append!(v, [5, 6])
end

# ╔═╡ 0a1011c2-c193-11ea-04a2-e71ff37bf002
md"a vector of vectors using an array comprehension"

# ╔═╡ d8b930b0-c193-11ea-0d53-63e3111f83a3
n_ = 20

# ╔═╡ c8dad7c0-c193-11ea-04a6-cf21953820d6
walkers = [trajectory(n_) for i ∈ 1:100]

# ╔═╡ 0f2b2db0-c194-11ea-06a1-6b2192a8e677
md"Use `end` to get the last element in a vector"

# ╔═╡ 61d6a940-c194-11ea-1d9f-4d6550479b45
[ traj_[end] for traj_ in walkers ] # iterating through walkers, a Vector of Vectors

# ╔═╡ 02bf22b0-c195-11ea-2d1b-ddebd9d7a069
md"Convert Vector of Vectors into a matrix:"

# ╔═╡ 13a7bfb0-c195-11ea-13d2-41437b78aae1
hcat(walkers...) # "splat" or like the "spread" operator in js

# ╔═╡ 188fbdb0-c196-11ea-16c2-0bbd3bd3c5f2
md"a more efficient way:"

# ╔═╡ b3074f30-c195-11ea-0025-af543bf613be
walkers_matrix = reduce(hcat, walkers)

# ╔═╡ 2a00013e-c196-11ea-1d45-3d8ff26e1e53
md"las element of the first row:"

# ╔═╡ f2221830-c195-11ea-09ff-2fb3e85bde34
walkers_matrix[1, end]

# ╔═╡ 5ed8c0a0-c196-11ea-2e0c-370732384c6b
md"last element of the first and second rows:"

# ╔═╡ 43f2c420-c196-11ea-3bfa-3320a1f4f92e
walkers_matrix[1:2, end]

# ╔═╡ 6e5598a0-c196-11ea-11d2-8f08d13fb07f
md"get the last element of all of the rows:"

# ╔═╡ 8b5d1592-c196-11ea-2a9f-17fbc031a2bd
walkers_matrix[1:end, end]

# ╔═╡ 9cbb09a0-c196-11ea-0975-c59ab691f596
md"also:"

# ╔═╡ a19bf82e-c196-11ea-024c-c7de6f0fbdb2
final_positions = walkers_matrix[:, end]

# ╔═╡ aed262a0-c196-11ea-310c-976af5163b2a
md"## Where are the walkers at last time step?"

# ╔═╡ e3dc4420-c196-11ea-37be-1388e05cdfc5
scatter(final_positions)

# ╔═╡ 8e8a1d20-c197-11ea-2241-7fed80778896
begin
	local n = 100

	local num_walkers = 10_000

	local walkers = [trajectory(n) for i ∈ 1:num_walkers]

	final_positions_ = [traj_[end] for traj_ ∈ walkers]
end

# ╔═╡ a46bcf22-c198-11ea-3496-cfdbdfcac559
scatter(final_positions_)

# ╔═╡ 527b1810-c198-11ea-3634-5f82562c0394
import Pkg; Pkg.add("StatsBase")

# ╔═╡ 76ce6910-c198-11ea-1a99-e33a29b2860f
counts = countmap(final_positions_)

# ╔═╡ c6170a90-c198-11ea-19c6-a377288360bd
scatter(counts)

# ╔═╡ Cell order:
# ╟─1c0741f0-c17a-11ea-3c6f-31f0c86dd08d
# ╟─58e6d60e-c17c-11ea-145e-ffdc97ab1f3f
# ╠═a93d6520-c17c-11ea-2cc7-5ff51f20df2b
# ╟─bfa68c10-c17c-11ea-281c-25e48e4ab81d
# ╠═ccef6d10-c17c-11ea-13e5-432b940ba726
# ╠═fea8c400-c17c-11ea-35fc-e98ef849c316
# ╟─9aff5fa0-c185-11ea-145a-c3230681bd4c
# ╠═0762ce10-c17d-11ea-1f2a-172e4ed0500a
# ╟─212df5a0-c186-11ea-0f5c-83f0dd5a1497
# ╠═8ce720ae-c185-11ea-380c-2df4848fb56a
# ╠═f52fc0a0-c185-11ea-2126-b10b11defe13
# ╟─891b2f20-c186-11ea-024f-372f0eef600c
# ╠═15990630-c186-11ea-2059-e1e301040dae
# ╠═62d228a2-c186-11ea-3ced-1db44003e269
# ╠═0b66e5a2-c187-11ea-36de-279c4d6aa180
# ╠═432008a0-c187-11ea-3606-159ae4ed8460
# ╟─131fcc00-c18a-11ea-09dc-0906c430d399
# ╟─1afa1280-c18c-11ea-1d1b-7f277d76e68c
# ╠═8cfcb9a2-c187-11ea-1129-ff3d55e25de4
# ╟─3b677390-c188-11ea-1079-01f262b1d76f
# ╠═4a816b10-c188-11ea-3935-1fc7371387ac
# ╠═643c5510-c188-11ea-2d15-fbf4939ae53c
# ╠═b417eaf0-c18c-11ea-3ae7-19f88eab6b9f
# ╟─eb35c2e0-c18d-11ea-3209-bd013b156820
# ╠═1ec51ce0-c18f-11ea-2166-6923cbcef091
# ╠═2c106ee0-c18f-11ea-21dc-db7d6706f5cb
# ╟─3ee28c60-c18f-11ea-3297-3d26c8d37627
# ╟─1ac71ed0-c190-11ea-2cfe-4d7cfd514800
# ╠═727357e0-c18e-11ea-24f4-81a80c9e5f36
# ╟─55e2ff60-c191-11ea-1b6b-b9a74d8fab5f
# ╟─60381bd0-c191-11ea-36a6-69333299f23b
# ╠═e1c193c0-c191-11ea-0c16-67d6d7611c26
# ╟─6fad5d40-c192-11ea-2573-01f0eb6c17ed
# ╠═f4e1b4c0-c192-11ea-3782-79f7c856869f
# ╟─0a1011c2-c193-11ea-04a2-e71ff37bf002
# ╠═d8b930b0-c193-11ea-0d53-63e3111f83a3
# ╠═c8dad7c0-c193-11ea-04a6-cf21953820d6
# ╟─0f2b2db0-c194-11ea-06a1-6b2192a8e677
# ╠═61d6a940-c194-11ea-1d9f-4d6550479b45
# ╟─02bf22b0-c195-11ea-2d1b-ddebd9d7a069
# ╠═13a7bfb0-c195-11ea-13d2-41437b78aae1
# ╟─188fbdb0-c196-11ea-16c2-0bbd3bd3c5f2
# ╠═b3074f30-c195-11ea-0025-af543bf613be
# ╟─2a00013e-c196-11ea-1d45-3d8ff26e1e53
# ╠═f2221830-c195-11ea-09ff-2fb3e85bde34
# ╟─5ed8c0a0-c196-11ea-2e0c-370732384c6b
# ╠═43f2c420-c196-11ea-3bfa-3320a1f4f92e
# ╟─6e5598a0-c196-11ea-11d2-8f08d13fb07f
# ╠═8b5d1592-c196-11ea-2a9f-17fbc031a2bd
# ╟─9cbb09a0-c196-11ea-0975-c59ab691f596
# ╠═a19bf82e-c196-11ea-024c-c7de6f0fbdb2
# ╟─aed262a0-c196-11ea-310c-976af5163b2a
# ╠═e3dc4420-c196-11ea-37be-1388e05cdfc5
# ╠═8e8a1d20-c197-11ea-2241-7fed80778896
# ╠═a46bcf22-c198-11ea-3496-cfdbdfcac559
# ╠═527b1810-c198-11ea-3634-5f82562c0394
# ╠═62a8f092-c198-11ea-2d26-632729a6bd7f
# ╠═76ce6910-c198-11ea-1a99-e33a29b2860f
# ╠═c6170a90-c198-11ea-19c6-a377288360bd


# Matern 5/2 ARD covariance Function

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

k(x,x') = σ²(1+√3*d/L + 5d²/3L²)exp(-√5*d/L), where d = |x-x'| and L = diag(ℓ₁,ℓ₂,...)
# Arguments:
* `ll::Vector{Float64}`: Log of the length scale ℓ
* `lσ::Float64`: Log of the signal standard deviation σ
""" ->
type Mat52Ard <: StationaryARD
    ℓ2::Vector{Float64}   # Log of Length scale
    σ2::Float64           # Log of signal std
    Mat52Ard(ll::Vector{Float64}, lσ::Float64) = new(exp(2.0*ll), exp(2.0*lσ))
end

function set_params!(mat::Mat52Ard, hyp::Vector{Float64})
    length(hyp) == num_params(mat) || throw(ArgumentError("Mat52 kernel only has $(num_params(mat)) parameters"))
    d = length(mat.ℓ2)
    mat.ℓ2 = exp(2.0*hyp[1:d])
    mat.σ2 = exp(2.0*hyp[d+1])
end

get_params(mat::Mat52Ard) = [log(mat.ℓ2)/2.0; log(mat.σ2)/2.0]
get_param_names(mat::Mat52Ard) = [get_param_names(mat.ℓ2, :ll); :lσ]
num_params(mat::Mat52Ard) = length(mat.ℓ2) + 1

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

function grad_kern(mat::Mat52Ard, x::Vector{Float64}, y::Vector{Float64})
    #r = distance(mat,x,y)
    wdiff = (x-y).^2./mat.ℓ2
    r = sqrt(sum(wdiff))
    exp_r = exp(-sqrt(5)*r)

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

    return [g1; g2]
end

function grad_stack!(stack::AbstractArray, mat::Mat52Ard, X::Matrix{Float64}, data::StationaryARDData)
    d = size(X,1)
    R = distance(mat,X)
    exp_R = exp(-sqrt(5)*R)

    stack[:,:,d+1] = 2.0*cov(mat, X)
    part = (5/3) * mat.σ2 .* exp_R .* (1.0 + sqrt(5)*R)
    broadcast!(*, view(stack, :, :, 1:d), (5/3)*mat.σ2, 1 + sqrt(5)*R, exp_R, data.dist_stack, reshape(1.0./mat.ℓ2, (1,1,d)))

    return stack
end


"""
VAR1SimPs(A,epsilon,T,x0=0.0)

Calculate impulse response function of a VAR(1) system
x(t) = A * x(t-1) +  epsilon(t), where x(t) is nx1

# Input
- `A::Matrix`:                  nxn VAR(1) matrix, see above
- `epsilon::Number or Vector`:  n-vector of shocks in inital period, or Txn matrix with shocks in all periods
- `T::Number`:                  scalar, last period to calculate for
- `x0::Number or Vector`:       n-vector with starting values, optional

# Output
- `xM::Matrix`:               Txn matrix, impulse response function

Paul.Soderlind@unisg.ch, to Julia Nov 2015

"""
function VAR1SimPs(A,epsilon,T,x0=0.0)

  n = size(A,1)

  isa(x0,Number)         && (x0      = fill(x0,n))           #if scalar
  isa(epsilon,Number)    && (epsilon = fill(epsilon,n))
  (length(epsilon) == n) && (epsilon = vcat(vec(epsilon)',zeros(T-1,n)))

  xM      = fill(NaN,(T,n))                        #to put results in
  xM[1,:] = A*vec(x0) + epsilon[1,:]
  for t = 2:T                                      #loop over time periods
    xM[t,:] = A*xM[t-1,:] + epsilon[t,:]
  end

  return xM

end
#-----------------------------------------------------------------------


module TestSystems
    using ..HomConBase

    export cyclic5, cyclic5Solutions, cyclic7, cyclic7Solutions

    """
        cyclical(iter, n)

    Cycles through an `iter` with n consecutive elements.

    ###Example:
        cyclical([1, 2, 3, 4], 2) == [[1, 2], [2, 3], [3, 4], [4, 1]]
    """
    function cyclical(iter, n)
        m = length(iter)
        if n == m
            return [collect(iter)]
        end
        values = collect(iter);
        res = Vector{Vector{eltype(values)}}()
        for k=0:m-1
            push!(res, [values[((k + i) % m) + 1] for i=0:n-1])
        end
        res
    end

    function monompoly(monomials::Vector{Vector{Int}}, totalvars)
        exps = zeros(Int, totalvars, length(monomials))
        for i in eachindex(monomials)
            for var in monomials[i]
                exps[var, i] = 1
            end
        end
        coeffs = ones(Int, length(monomials))
        Poly(exps, coeffs)
    end

    function cyclical_polys(n)
        F = map(k -> monompoly(cyclical(1:n, k), n), 1:n-1)
        # now we have to construct x₁x₂x₃...x_n - 1
        exps = zeros(Int, n, 2)
        exps[:,1] = 1
        push!(F, Poly(exps, [1, -1]))
        PolySystem(F, [Symbol("z$i") for i=1:n])
    end

    """
        cyclic5()

    Cyclic5 example problem folowing [^1]

    [^1]: A faster way to count the solutions of inhomogeneous systems of algebraic equations,
    with applications to cyclic n-roots
    """
    function cyclic5()
        cyclical_polys(5)
    end

    """
        cyclic5Solutions()

    Solutions to the Cyclic5 example.
    """
    function cyclic5Solutions()
        # (1,w,w^{2k},w^{3k},w^{4k}) with w = exp(2πi/5) permuted cyclical, 1 ≦ k ≦ 4
        w = exp(2.0*π*im/5.0)
        classic_sols = vec([ map(i -> w^(i*k), perm) for k in 1:4, perm in cyclical(0:4, 5)])

        #solution of ɛ^2+3ɛ+1=0
        ɛ = -1.5 - 0.5 * √5.0 + 0*im
        extended_sols = vec([ rev ? reverse!(w^k * perm) : w^k * perm for perm in cyclical([ɛ, 1/ɛ, 1, 1, 1], 5),
                                                                            k in 0:4,
                                                                            rev in [false, true]])
        vcat(classic_sols, extended_sols)
    end
    """
        cyclic7()

    Cyclic7 problem folowing [^1]

    [^1]: A faster way to count the solutions of inhomogeneous systems of algebraic equations,
    with applications to cyclic n-roots
    """
    function cyclic7()
        cyclical_polys(7)
    end

    """
        cyclic7Solutions()

    Solutions to the Cyclic7 example.
    """
    function cyclic7Solutions()
        # (1,w,w^{2k},w^{3k},w^{4k},w^{5k},w^{6k}) with w = exp(2πi/7) permuted cyclical, 1 ≦ k ≦ 6
        w = exp(2.0*π*im/7.0)
        classic_sols = vec([ map(i -> w^(i*k), perm) for k in 1:6, perm in cyclical(0:6, 7)])

        #solution of ɛ^2+5ɛ+1=0
        ɛ = -2.5 - 0.5 * √21.0 + 0*im
        extended_sols = vec([ rev ? reverse!(w^k * perm) : w^k * perm for perm in cyclical([ɛ, 1/ɛ, 1, 1, 1, 1, 1], 7),
                                                                            k in 0:6,
                                                                            rev in [false, true]])

        # index two solutions
        d = 0.25*im*(√7 + 3*im)
        index_two = vec([ rev ? reverse!(w^k * perm) : w^k * perm for perm in cyclical([d, conj(d), d, 1, conj(d), 1, 1], 7),
                                                                            k in 0:6,
                                                                            rev in [false, true]])
        index_two_conj = vec([ rev ? reverse!(w^k * perm) : w^k * perm for perm in cyclical([conj(d), d, conj(d), 1, d, 1, 1], 7),
                                                                            k in 0:6,
                                                                            rev in [false, true]])

        # index three solutions
        a = 2.738895317095 + 0im
        b = 3.436680125767 + 0im
        c = -0.1298393513967 + 0im
        sol_1 = [a, b, c, 1, 1/c, 1/b, 1/a]
        sol_2 = [a*b,c,1/(b*c),1,b*c,1/c, 1/(a*b)]
        sol_3 = [a*b*c,1/(b*c),b,1,1/b,b*c,1/(a*b*c)]

        perms = [cyclical(sol_1, 7); cyclical(sol_2, 7); cyclical(sol_3, 7)]
        index_three_first = vec([ rev ? reverse!(w^k * perm) : w^k * perm for perm in perms, k in 0:6, rev in [false, true]])

        a = exp(4.3128389787245im)
        b = exp(1.356227956787im)
        c = exp(1.900668281165im)
        sol_1 = [a, b, c, 1, 1/c, 1/b, 1/a]
        sol_2 = [a*b,c,1/(b*c),1,b*c,1/c, 1/(a*b)]
        sol_3 = [a*b*c,1/(b*c),b,1,1/b,b*c,1/(a*b*c)]
        perms = [cyclical(sol_1, 7); cyclical(sol_2, 7); cyclical(sol_3, 7)]
        index_three_second = vec([ rev ? reverse!(w^k * perm) : w^k * perm for perm in perms, k in 0:6, rev in [false, true]])

        vcat(classic_sols, extended_sols, index_two, index_two_conj, index_three_first, index_three_second)
    end
end


# Using Convex.jl to test if the probability distributions in Homework problem 1
# Sets up the optimization problem for the Cut Inflation for the triangle causal
# model with the distribution:
#
# P(a, b, c) = if (a + b + c == 1) => 1/3
#              else             => 0

using Convex
using GLPK
using Inflation

# the probability matrix is set up as
# P_{A1,B1,C1} = sum_{A2,B2,C2}P     --> 'yxyxyx'
# P_{A1,B2,C2} = sum_{A2,B1,C1}P     --> 'yxxyxy'
# P_{A2,B1,C2} = sum_{A1,B2,C1}P     --> 'xyyxxy'
# P_{A2,B2,C1} = sum_{A1,B1,C2}P     --> 'xyxyyx'
# P_{A2,B2,C2} = sum_{A1,B1,C1}P     --> 'xyxyxy'
infl_strings = ["yxyxyx", "yxxyxy", "xyyxxy", "xyxyyx", "xyxyxy"]

# creating the M-matrix and b-array to be used in the Convex.jl optimization
M = get_M_matrix(infl_strings)
b = get_b_array(infl_strings, pfunc2, spiral)

# the variables that get optimized by Convex.jl
v = get_v_array_Variable(infl_strings)

# set up the objective and constraints
objective   = 0
constraints = [v >= 0, M*v == b]

# create the problem and solve
problem = Convex.maximize(objective, constraints)
solve!(problem, GLPK.Optimizer, verbose=false)

# return the results
if problem.status == Convex.MathOptInterface.INFEASIBLE
	println("Solution is infeasible.")
else
	println(problem.optval)
	println(v.value)
end


using Distributions
function discrete_normal(n,mu,sigma,width)

# creates equally spaced approximation to normal distribution
# n is number of points
# mu is mean
# sigma is standard deviation
# width is the multiple of stand deviation for the width of the grid
# f is the error in the approximation
# x gives the location of the points
# p is probabilities

    x = collect(range(mu-width*sigma,mu+width*sigma,length=n))
    if n==2
        p = 0.5.*ones(n)
    elseif n>2
        p  = zeros(n)
        p[1] = cdf(Normal(mu,sigma), x[1] + 0.5*(x[2]-x[1]))
        for i = 2:n-1
            p[i] = cdf(Normal(mu,sigma), x[i] + 0.5*(x[i+1]-x[i])) - cdf(Normal(mu,sigma), x[i] - 0.5*(x[i]-x[i-1]))
        end
        p[n] = 1 - sum(p[1:n-1])
    end

    Ex = x'*p
    SDx = sqrt((x.^2)'*p .- Ex.^2)

    f = SDx.-sigma

    return f,x,p

end


function bin_dist(n, k, p, q)
    return binomial(n, k) * p^k * q^(n-k)
end

function f(n, k, p, q)
    mass = 0
    for i = 0:k
        mass += bin_dist(n, i, p, q)
    end
    return mass
end

function g(n, k, p, q)
    mass = 0
    for i = k:n
        mass += bin_dist(n, i, p, q)
    end
    return mass
end

str = readline()
strvec = split(str, ' ')
intvec = map(x -> parse(Int64, x), strvec)
a = intvec[1]
n = intvec[2]

p = a / 100
q = 1 - p

println(round(f(n, 2, p, q), 3))
println(round(g(n, 2, p, q), 3))

Tutorials/10 Days of Statistics/Day 4/binomial_distribution2.jl


# ---
# title: 69. Sqrt(x)
# id: problem69
# author: Indigo
# date: 2022-03-03
# difficulty: Easy
# categories: Math, Binary Search
# link: <https://leetcode.com/problems/sqrtx/description/>
# hidden: true
# ---
#
# Given a non-negative integer `x`, compute and return _the square root of_ `x`.
#
# Since the return type is an integer, the decimal digits are **truncated** ,
# and only **the integer part** of the result is returned.
#
#
#
# **Example 1:**
#
#
#
#     Input: x = 4
#     Output: 2
#
#
# **Example 2:**
#
#
#
#     Input: x = 8
#     Output: 2
#     Explanation: The square root of 8 is 2.82842..., and since the decimal part is truncated, 2 is returned.
#
#
#
# **Constraints:**
#
#   * `0 <= x <= 231 - 1`
#
#
## @lc code=start
using LeetCode

my_sqrt = isqrt
function mysqrt(x::Int)::Int
    x < 0 && throw(DomainError(x))
    left, right = 0, x
    while left <= right
        mid = left + ((right - left) >> 1)
        if mid ^ 2 > x
            right = mid - 1
        else
            left = mid + 1
        end
    end
    right
end

## @lc code=end


mutable struct FixedStepOdeOptions
    step_size::Float64
end
function FixedStepOdeOptions(; step_size = 1e-2)
    FixedStepOdeOptions(step_size)
end

function ode_solver_explicit_euler(fcn, time_interval::TimeInterval, initial_values::Vector{T}, options::FixedStepOdeOptions) where {T}

    step_size = options.step_size

    t     = time_interval.t_start
    q     = initial_values

    ode_res = init_ode_result(q)

    while true

        dqdt  = fcn(t, q)
        q_new = q + step_size * dqdt

        t = t + step_size
        q = q_new

        push!(ode_res, t, q)

        if t >= time_interval.t_end
            break
        end

    end

    return ode_res
end

function ode_solver_semiimplicit_euler(fcn, time_interval::TimeInterval, initial_values::Vector{T}, options::FixedStepOdeOptions) where {T}

    step_size = options.step_size

    t     = time_interval.t_start
    q     = initial_values

    ode_res = init_ode_result(q)

    n_dim = length(initial_values)
    id = Matrix{T}( one(T) * I , n_dim, n_dim )

    while true

        t_new = t + step_size
        jac = numder.jacobian_fw( q->fcn(t_new, q), q)

        mat = id - step_size * jac
        dq = step_size .*( mat \ fcn(t_new, q) )

        t = t_new
        q = q + dq

        push!(ode_res, t, q)

        if t >= time_interval.t_end
            break
        end

    end

    return ode_res

end

function ode_solver_implicit_euler(fcn, time_interval::TimeInterval, initial_values::Vector{T}, options::FixedStepOdeOptions) where {T}

    step_size = options.step_size

    t     = time_interval.t_start
    q     = initial_values

    ode_res = init_ode_result(q)

    n_dim = length(initial_values)
    # id = Matrix{T}( one(T) * I , n_dim, n_dim )

    while true

        t_new = t + step_size
        nonlin_fcn(y) = step_size .* fcn(t_new, y) - y + q

        q_try = q + step_size * fcn(t_new, q)

        (q_new, nonlineq_iter) = nonlineq.broyden( nonlin_fcn, q_try )

        t = t_new
        q = q_new

        push!(ode_res, t, q)

        if t >= time_interval.t_end
            break
        end

    end

    return ode_res

end

function ode_solver_rk4(
    fcn,
    time_interval::TimeInterval,
    initial_values::Vector{T},
    options::FixedStepOdeOptions) where {T}

end


# =========================================================================================================================
# Theoretical error bounds for Carleman linearization
#
# References:
#
# - [1] Forets, Marcelo, and Amaury Pouly. "Explicit error bounds for carleman linearization."
#       arXiv preprint arXiv:1711.02552 (2017).
#
# - [2] Liu, J. P., Kolden, H. Ø., Krovi, H. K., Loureiro, N. F., Trivisa, K., & Childs,7
#       A. M. (2021). "Efficient quantum algorithm for dissipative nonlinear differential equations."
#       Proceedings of the National Academy of Sciences, 118(35). arXiv preprint arXiv:2011.03185.
#
# =========================================================================================================================

# --- Error bounds using a priori estimate from [1] ---

# See Theorem 4.2 in [1]. This is a bound based on an a priori estimate
# of the norm of the exact solution x(t).
# These bounds use the supremum norm (p = Inf).
function error_bound_apriori(α, F₁, F₂; N)
    nF₂ = opnorm(F₂, Inf)
    μF₁ = logarithmic_norm(F₁, Inf)

    β = α * nF₂ / μF₁
    ε = t -> α * β^N * (exp(μF₁ * t) - 1)^N
    return ε
end

# See Theorem 4.2 in [1]
function convergence_radius_apriori(α, F₁, F₂; N)
    nF₂ = opnorm(F₂, Inf)
    μF₁ = logarithmic_norm(F₁, Inf)

    if μF₁ < 0
        return Inf
    end
    β = α * F₂ / μF₁
    T = (1/μF₁) * log(1 + 1/β)
    return T
end

# --- Error bounds using power series method from [1] ---

# See Theorem 4.3 in [1], which uses the power series method.
function error_bound_pseries(x₀, F₁, F₂; N)
    nx₀ = norm(x₀, Inf)
    nF₁ = opnorm(F₁, Inf)
    nF₂ = opnorm(F₂, Inf)
    β₀ = nx₀ * nF₂ / nF₁

    ε = t -> nx₀ * exp(nF₁ * t) / (1 - β₀ * (exp(nF₁ * t) - 1)) * (β₀ * (exp(nF₁ * t) - 1))^N
    return ε
end

# See Theorem 4.3 in [1].
function convergence_radius_pseries(x₀, F₁, F₂; N)
    nx₀ = norm(x₀, Inf)
    nF₁ = opnorm(F₁, Inf)
    nF₂ = opnorm(F₂, Inf)
    β₀ = nx₀ * nF₂ / nF₁

    T = (1/nF₁) * log(1 + 1/β₀)
    return T
end

# --- Error bounds using spectral abscissa from [2] ---

# See Definition (2.2) in [2]. These bounds use the spectral norm (p = 2)
function _error_bound_specabs_R(x₀, F₁, F₂; check=true)
    nx₀ = norm(x₀, 2)
    nF₂ = opnorm(F₂, 2)

    # compute eigenvalues and sort them by increasing real part
    λ = eigvals(F₁, sortby=real)
    λ₁ = last(λ)
    Re_λ₁ = real(λ₁)
    if check
        @assert Re_λ₁ <= 0 "expected Re(λ₁) ≤ 0, got $Re_λ₁"
    end
    R = nx₀ * nF₂ / abs(Re_λ₁)
    return (R, Re_λ₁)
end

# See Lemma 2 in [2]
function error_bound_specabs(x₀, F₁, F₂; N, check=true)
    (R, Re_λ₁) = _error_bound_specabs_R(x₀, F₁, F₂; check=check)
    if check
        @assert R < 1 "expected R < 1, got R = $R; try scaling the ODE"
    end

    nx₀ = norm(x₀, 2)
    if iszero(Re_λ₁)
        nF₂ = opnorm(F₂, 2)
        ε = t -> nx₀ * (nx₀ * nF₂ * t)^N
    else
        ε = t -> nx₀ * R^N * (1 - exp(Re_λ₁ * t))^N
    end
    return ε
end

# See Lemma 2 in [2]
function convergence_radius_specabs(x₀, F₁, F₂; check=true)
    (R, Re_λ₁) = _error_bound_specabs_R(x₀, F₁, F₂; check=check)

    if Re_λ₁ < 0
        T = Inf
    elseif iszero(Re_λ₁)
        nx₀ = norm(x₀, 2)
        nF₂ = opnorm(F₂, 2)
        β = nx₀ * nF₂
        T = 1/β
    else
        throw(ArgumentError("expected spectral abscissa to be negative or zero, got $Re_λ₁"))
    end
    return T
end


using LinearAlgebra
using ExtendableSparse

mutable struct RadialBasis{F,Q,X,Y,L,U,C,S,D} <: AbstractSurrogate
    phi::F
    dim_poly::Q
    x::X
    y::Y
    lb::L
    ub::U
    coeff::C
    scale_factor::S
    sparse::D
end

mutable struct RadialFunction{Q,P}
    q::Q # degree of polynomial
    phi::P
end

linearRadial = RadialFunction(0,z->norm(z))

cubicRadial = RadialFunction(1,z->norm(z)^3)
multiquadricRadial = RadialFunction(1,z->sqrt(norm(z)^2+1))

thinplateRadial = RadialFunction(2, z->begin
    result = norm(z)^2 * log(norm(z))
    ifelse(iszero(z), zero(result), result)
end)

"""
    RadialBasis(x,y,lb::Number,ub::Number; rad::RadialFunction = linearRadial,scale::Real=1.0)

Constructor for RadialBasis surrogate.
"""
function RadialBasis(x, y, lb::Number, ub::Number; rad::RadialFunction=linearRadial, scale_factor::Real=1.0, sparse = false)
    q = rad.q
    phi = rad.phi
    coeff = _calc_coeffs(x, y, lb, ub, phi, q,scale_factor, sparse)
    return RadialBasis(phi, q, x, y, lb, ub, coeff,scale_factor,sparse)
end

"""
RadialBasis(x,y,lb,ub,rad::RadialFunction, scale_factor::Float = 1.0)

Constructor for RadialBasis surrogate
"""
function RadialBasis(x, y, lb, ub; rad::RadialFunction = linearRadial, scale_factor::Real=1.0, sparse = false)
    q = rad.q
    phi = rad.phi
    coeff = _calc_coeffs(x, y, lb, ub, phi, q, scale_factor, sparse)
    return RadialBasis(phi, q, x, y, lb, ub, coeff,scale_factor, sparse)
end

function _calc_coeffs(x, y, lb, ub, phi, q, scale_factor, sparse)
    nd = length(first(x))
    num_poly_terms = binomial(q + nd, q)

    D = _construct_rbf_interp_matrix(x, first(x), lb, ub, phi, q, scale_factor, sparse)
    Y = _construct_rbf_y_matrix(y, first(y), length(y) + num_poly_terms)

    coeff = D \ Y
    return coeff
end

function _construct_rbf_interp_matrix(x, x_el::Number, lb, ub, phi, q, scale_factor, sparse)
    n = length(x)

    num_poly_terms = binomial(q + 1, q)
    m = n + num_poly_terms

    if sparse
        D = ExtendableSparseMatrix{eltype(x_el),Int}(m,m)
    else
        D = zeros(eltype(x_el), m, m)
    end
    @inbounds for i = 1:n
        for j = 1:n
            D[i,j] = phi( (x[i] .- x[j]) ./ scale_factor )
        end
        if i <= n
            for k = 1:num_poly_terms
                    D[i,n+k] = _scaled_chebyshev(x[i], k-1, lb, ub)
            end
        end
    end
    D_sym = Symmetric(D, :U)
    return D_sym
end

function _construct_rbf_interp_matrix(x, x_el, lb, ub, phi, q, scale_factor,sparse)
    n = length(x)
    nd = length(x_el)

    num_poly_terms = binomial(q + nd, q)
    m = n + num_poly_terms

    if sparse
        D = ExtendableSparseMatrix{eltype(x_el),Int}(m,m)
    else
        D = zeros(eltype(x_el), m, m)
    end
    @inbounds for i = 1:n
        for j = 1:n
            D[i,j] = phi( (x[i] .- x[j]) ./ scale_factor)
        end
        if i < n + 1
            for k = 1:num_poly_terms
                D[i,n+k] = multivar_poly_basis(x[i], k-1, nd, q)
            end
        end
    end
    D_sym = Symmetric(D, :U)
    return D_sym
end

_construct_rbf_y_matrix(y, y_el::Number, m) = [i <= length(y) ? y[i] : zero(y_el) for i = 1:m]
_construct_rbf_y_matrix(y, y_el, m) = [i <= length(y) ? y[i][j] : zero(first(y_el)) for i=1:m, j=1:length(y_el)]

using Zygote: @nograd

function _make_combination(n, d, ix)
    exponents_combinations = [
        e
        for e
        in collect(
            Iterators.product(
                Iterators.repeated(0:n, d)...
            )
        )[:]
        if sum(e) <= n
    ]

    return exponents_combinations[ix + 1]
end
# TODO: Is this correct? Do we ever want to differentiate w.r.t n, d, or ix?
# By using @nograd we force the gradient to be 1 for n, d, ix
@nograd _make_combination

"""
    multivar_poly_basis(x, ix, d, n)

Evaluates in `x` the `ix`-th element of the multivariate polynomial basis of maximum
degree `n` and `d` dimensions.

Time complexity: `(n+1)^d.`

# Example
For n=2, d=2 the multivariate polynomial basis is
````
1,
x,y
x^2,y^2,xy
````
Therefore the 3rd (ix=3) element is `y` .
Therefore when x=(13,43) and ix=3 this function will return 43.
"""
multivar_poly_basis(x, ix, d, n) = prod(
    a^d
    for (a, d)
    in zip(x, _make_combination(n, d, ix))
)

"""
Calculates current estimate of value 'val' with respect to the RadialBasis object.
"""
function (rad::RadialBasis)(val)
    approx = _approx_rbf(val, rad)
    return _match_container(approx, first(rad.y))
end

function _approx_rbf(val::Number, rad)
    n = length(rad.x)
    q = rad.dim_poly
    num_poly_terms = binomial(q + 1, q)
    lb = rad.lb
    ub = rad.ub
    approx = zero(rad.coeff[1, :])
    for i = 1:n
        approx += rad.coeff[i, :] * rad.phi( (val .- rad.x[i]) / rad.scale_factor)
    end
    for k = 1:num_poly_terms
        approx += rad.coeff[n+k, :] * _scaled_chebyshev(val, k-1, lb, ub)
    end
    return approx
end
function _approx_rbf(val, rad)
    n = length(rad.x)
    d = length(rad.x[1])
    q = rad.dim_poly
    num_poly_terms = binomial(q + d, q)
    lb = rad.lb
    ub = rad.ub
    sum_half_diameter = sum((ub[k]-lb[k])/2 for k = 1:d)
    mean_half_diameter = sum_half_diameter/d
    central_point = _center_bounds(first(rad.x), lb, ub)

    approx = zero(rad.coeff[1, :])
    @views approx += sum( rad.coeff[i, :] * rad.phi( (val .- rad.x[i]) ./rad.scale_factor) for i = 1:n)
    for k = 1:num_poly_terms
        @views approx += rad.coeff[n+k, :] .* multivar_poly_basis(val, k-1, d, q)
    end
    return approx
end

_scaled_chebyshev(x, k, lb, ub) = cos(k*acos(-1 + 2*(x-lb)/(ub-lb)))
_center_bounds(x::Tuple, lb, ub) = ntuple(i -> (ub[i] - lb[i])/2, length(x))
_center_bounds(x, lb, ub) = (ub .- lb) ./ 2

"""
    add_point!(rad::RadialBasis,new_x,new_y)

Add new samples x and y and update the coefficients. Return the new object radial.
"""
function add_point!(rad::RadialBasis,new_x,new_y)
    if (length(new_x) == 1 && length(new_x[1]) == 1) || ( length(new_x) > 1 && length(new_x[1]) == 1 && length(rad.lb)>1)
        push!(rad.x,new_x)
        push!(rad.y,new_y)
    else
        append!(rad.x,new_x)
        append!(rad.y,new_y)
    end
    rad.coeff = _calc_coeffs(rad.x,rad.y,rad.lb,rad.ub,rad.phi,rad.dim_poly, rad.scale_factor, rad.sparse)
    nothing
end


### A Pluto.jl notebook ###
# v0.14.0

using Markdown
using InteractiveUtils

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

# ╔═╡ 430d5134-92af-11eb-1fcc-d71f5216d62a
#Set up the working environment
begin
	using Pkg
	Pkg.activate("..")
	using Plots, PlutoUI,LaTeXStrings, DifferentialEquations
end

# ╔═╡ 5adc7e34-92af-11eb-2cfe-67899a4a0291
md"""
## Attempting a more realistic model

Next, we'll try to make the model more realistic by including **reinfection**, **vaccination** and **mortalitiy**.

##### Mortality
We'll only consider deaths of infected people at a mortality rate $m$: $$I \xrightarrow{m} D$$

##### Vaccination & Reinfection
To account for vaccination, we'll assume that
1. The rate of infection is different for vaccinated people $V$ and susceptible people $S$, represented by a *vaccination immunity factor*, $i_v < 1$

$$S \xrightarrow{r_s} I \xleftarrow{r_s \cdot i_v} V$$

2. Likewise, susceptible people $S$ and recovered people $R$ become infected at different rates, represented by an *recovered immunity factor*, $i_r < 1$ that lowers the infection risk

$$S \xrightarrow{r_s} I \xleftarrow{r_s \cdot i_r} R$$

3. We assume that recovered people $R$ and susceptible people $S$ are vaccinated at the same rate $v$

$$S \xrightarrow{v} V \xleftarrow{v} R$$

Taken together, this yields the following set of differential equations:

$$\begin{align}
	\frac{dS}{dt} = -rSI -vS\\
	\frac{dV}{dt} = vS + vR - i_v r VI\\
	\frac{dD}{dt} = mI\\
	\frac{dR}{dt} = aI - i_r rRI - vR\\
	\frac{dI}{dt} = rSI + i_r rRI + i_v r VI - mI - aI
\end{align}$$

"""

# ╔═╡ 7a19b238-92af-11eb-30fd-b189d06a0882
md""" ## Simulation
*Change the sliders to adjust the corresponding values*:

| parameter                   | value                     |
|-----------------------------|---------------------------|
| transmission rate $r$       | $(@bind r2 Slider(0:0.005:1, default=0.125, show_value=true)) |
| recovery rate $a$           | $(@bind a2 Slider(0:0.001:0.25, default=0.1, show_value=true))|
| rate of vaccination $v$     | $(@bind v2 Slider(0:0.0001:0.01, default=0.0008, show_value=true)) |
| mortality $m$               | $(@bind m2 Slider(0:0.0001:0.05, default=0.0006, show_value=true)) |
| % immunity recovered        | $(@bind ir Slider(0:1.0:100.0, default=60, show_value=true)) |
| % immunity vaccinated       | $(@bind iv Slider(0:1.0:100.0, default=90, show_value=true)) |
|  $I(0)$                     | $(@bind I02 Slider(0.0001:0.001:0.1, default=0.004, show_value=true)) |
|  $R(0)$                     | $(@bind R02 Slider(0:0.001:0.1, default=0.055,show_value=true)) |
| $V(0)$                      | $(@bind V02 Slider(0:0.001:0.1, default=0.048,show_value=true)) |
| $D(0)$                      | $(@bind D02 Slider(0:0.001:0.1, default=0.001,show_value=true)) |
| timespan $t_{max}$          | $(@bind tmax2 Slider(10.0:10.0:2000.0, default=100.0,show_value=true)) days |
| zoom y-axis                 | $(@bind ymax2 Slider(0.001:0.001:1.0, default=1.0,show_value=true)) |
| Population size             | $(@bind popsize NumberField(500000:10000000000, default=9044650)) |
"""

# ╔═╡ 8f867f9a-92af-11eb-3671-43904aa74d8b
function getres(solution, i, t, p)
	v = solution(t)[i]*p
	return round(maximum([0.0 v]))
end

# ╔═╡ 96d8a5c2-92af-11eb-1d3f-9bdd035eafec
#S, I, R, V, D
function sirvd!(du,u,p,t)
	du[1] = -u[1] * ( p[1] * u[2] + p[3])
	du[2] = u[2] * (p[1] * (u[1] + p[5] * u[4] + p[4] * u[3]) - p[6] - p[2])
	du[3] = p[2] * u[2] - u[3] * ( p[1] * p[4] *  u[2] + p[3])
	du[4] = p[3] * (u[1] + u[3]) - u[4] * u[2] * p[1] * p[5]
	du[5] = p[6] * u[2]
end

# ╔═╡ 9c0dfe82-92af-11eb-3192-49c2deac6e9b
p2 = [r2, a2, v2, (100 - ir)/100, (100 - iv)/100, m2]

# ╔═╡ a1e4ca5e-92af-11eb-1c40-410a8d3abaca
u02 = [1-I02-R02-V02-D02, I02, R02, V02, D02]

# ╔═╡ ada6253e-92af-11eb-0e32-d9760ab88271
tspan2 = (0.0, 2000)

# ╔═╡ b08f47b2-92af-11eb-367b-979567aebb4d
prob2 = ODEProblem(sirvd!,u02,tspan2,p2)

# ╔═╡ 88b411f0-92af-11eb-2abf-151b2bfb313b
begin
	sol2 = solve(prob2, dt=1.0, dense = true)
	infected = sol2[2,:]*popsize
	dead = infected + sol2[5,:]*popsize
	recovered = dead + sol2[3,:]*popsize
	vaccinated = recovered + sol2[4,:]*popsize
	susceptible = vaccinated + sol2[1,:]*popsize
	icu = infected*0.015
	icu_cap = [1000 for t in sol2.t]

	sim_stackarea = plot(sol2.t, [susceptible, vaccinated, recovered, dead, infected, icu],
				fill = true,
				color = ["gray89" "skyblue3" "paleturquoise3" "gray39" "lightcoral" "maroon"],
				fillalpha = 1,
				label = ["S(t)" "V(t)" "R(t)" "D(t)" "I(t)" L"$I(t)_{icu}$"],
				xlabel = "t / days",
				title = "Timecourse (stacked area)",
				xlims = (0, tmax2),
				ylims = (0, ymax2*popsize))

	sim_timecourse = plot(sol2,
				color = ["gray89" "lightcoral" "paleturquoise3" "skyblue3" "gray39"],
				label = ["S(t)" "I(t)" "R(t)" "V(t)" "D(t)"],
				xlabel = "t / days",
				title = "Timecourse",
				xlims = (0, tmax2),
				ylims = (0, ymax2))

	sim_rate_tc = plot(0:tmax2, sol2(0:tmax2, Val{1}, idxs=2),
				color = ["lightcoral"],
				label = [L"\frac{dI}{dt}"],
				xlabel = "t / days",
				title = "New Infections",
				xlims = (0, tmax2))

	timecourses = plot(sim_timecourse, sim_rate_tc, layout = (2,1))

	plot(sim_stackarea, timecourses, layout = (1, 2),fontfamily="Computer Modern")
end

# ╔═╡ 7ec0712a-92af-11eb-2144-15d33af27c3d
md"""
| t (days) | I(t)                               | R(t)                               | V(t)                               | D(t)                               |
|----------|------------------------------------|------------------------------------|------------------------------------|------------------------------------|
| 0        | $(getres(sol2, 2, 0, popsize))     | $(getres(sol2, 3, 0, popsize))     | $(getres(sol2, 4, 0, popsize))     | $(getres(sol2, 5, 0, popsize))     |
| $tmax2   | $(getres(sol2, 2, tmax2, popsize)) | $(getres(sol2, 3, tmax2, popsize)) | $(getres(sol2, 4, tmax2, popsize)) | $(getres(sol2, 5, tmax2, popsize)) |

Base reprodution rate = $(round((r2*(1-I02-R02-V02-D02)/a2), sigdigits = 3)),
Infectious period = $(round(1/a2, sigdigits = 3))

"""

# ╔═╡ 230527fa-92b3-11eb-308f-ffab3777029c
sol2(0:tmax2, Val{1}, idxs=2)

# ╔═╡ Cell order:
# ╟─430d5134-92af-11eb-1fcc-d71f5216d62a
# ╟─5adc7e34-92af-11eb-2cfe-67899a4a0291
# ╟─7a19b238-92af-11eb-30fd-b189d06a0882
# ╟─7ec0712a-92af-11eb-2144-15d33af27c3d
# ╟─88b411f0-92af-11eb-2abf-151b2bfb313b
# ╟─230527fa-92b3-11eb-308f-ffab3777029c
# ╟─8f867f9a-92af-11eb-3671-43904aa74d8b
# ╟─96d8a5c2-92af-11eb-1d3f-9bdd035eafec
# ╟─9c0dfe82-92af-11eb-3192-49c2deac6e9b
# ╟─a1e4ca5e-92af-11eb-1c40-410a8d3abaca
# ╟─ada6253e-92af-11eb-0e32-d9760ab88271
# ╟─b08f47b2-92af-11eb-367b-979567aebb4d


__precompile__()

module Napier

export loge, logn, ln

function loge(n::Real, li::Real, ls::Real)
  if abs(li-ls) <= 0.000001
    (li+ls)/2.0
  elseif (exp(li)-n)*(exp((li+ls)/2.0)-n) < 0
    loge(n,li,(li+ls)/2.0)
  else
    loge(n,(li+ls)/2.0,ls)
  end
end

function logn(n::Real)
  if n <= 0
    throw(DomainError())
  elseif n < 1
    loge(n, 0, -n-80)
  elseif n == 1
    0
  else
    loge(n, 0, n)
  end
end

const ln = logn

end # module


function VorticityBoundaryConditions!(ω,ψ,Δx,Δy,un,us,ve,vw)
  ω[:,end] .= 2*((ψ[:,end]  - ψ[:,end-1] )/(Δx^2) .- ve/Δx)
  ω[:,1] .= 2*((ψ[:,1]  - ψ[:,2]   )/(Δx^2) .- vw/Δx)
  ω[end,:] .= 2*((ψ[end,:]  - ψ[end-1,:] )/(Δy^2) .+ us/Δy)
  ω[1,:] .= 2*((ψ[1,:]  - ψ[2,:]   )/(Δy^2) .+ un/Δy)
end
function GaussSiedel!(ϕ,Ap,An,As,Ae,Aw,Rp,res; λ=1, maxiter = 1000)
  normRes = 1
  k = 0
  Ny,Nx = size(ϕ)
  while normRes >= 1e-8 && k < maxiter
    k += 1
    for i in 2:Ny-1
      for j in 2:Nx-1
        ϕP = ϕ[i,j]
        ϕE = ϕ[i+0,j+1]
        ϕW = ϕ[i+0,j-1]
        ϕN = ϕ[i-1,j+0]
        ϕS = ϕ[i+1,j+0]
        res[i,j] = Rp[i,j] - (Ap*ϕP
          + An*ϕN
          + As*ϕS
          + Ae*ϕE
          + Aw*ϕW)
        Δϕ = res[i,j]/Ap
        ϕ[i,j] = λ*(ϕ[i,j] + Δϕ) + (1-λ)*ϕ[i,j]
      end
    end
    normRes = norm(res)
  end
  return k
end
function LinearSolve!(A,x,b)
  # Solves the equation Ax = b assuming zero Dirichlet BCs everywhere
  Ny,Nx = size(b)
  Ny,Nx = Ny-2, Nx-2
  x_int = x[2:end-1,2:end-1]
  b_int = b[2:end-1,2:end-1]
  b_vec = reshape(b_int,Ny*Nx)
  # x_int = A\b_vec
  x_vec = reshape(x_int,Ny*Nx)
  cg!(x_vec,A,b_vec, log = true)
  x[2:end-1,2:end-1] .= reshape(x_int,(Ny,Nx))
end
function BuildPoissonMatrix(Ny,Nx,Δx,Δy)
  # This function returns a (Ny*Nx) × (Ny*Nx) matrix in the form of
  # a sparse array, corresponding to the discrete 2D Laplacian operator.
  Ny = Ny-2
  Nx = Nx-2

  Isx = [1:Ny; 1:Ny-1; 2:Ny]
  Jsx = [1:Ny; 2:Ny; 1:Ny-1]

  Isy = [1:Nx; 1:Nx-1; 2:Nx]
  Jsy = [1:Nx; 2:Nx; 1:Nx-1]

  Vsx = [fill(-2,Ny); fill(1, 2Ny-2)]
  Vsy = [fill(-2,Nx); fill(1, 2Nx-2)]
  D²x = sparse(Isx, Jsx, Vsx)
  D²y = sparse(Isy, Jsy, Vsy)
  # D_xx = 1/(Δx^2) .* kron(sparse(I,Nx,Nx), D²x)
  # D_yy = 1/(Δy^2) .* kron(D²y, sparse(I,Ny,Ny))
  D_yy = 1/(Δy^2) .* kron(sparse(I,Nx,Nx), D²x)
  D_xx = 1/(Δx^2) .* kron(D²y, sparse(I,Ny,Ny))
  Lap = D_xx + D_yy
end
function BuildAdvectionDiffusionCoefficients(Re,Δt,Δx,Δy)
  # Time-derivative
  ap = 1/Δt
  # Diffusion
  ap += 2/(Re*Δx^2) + 2/(Re*Δy^2)
  an = -1/(Re*Δy^2)
  aw = -1/(Re*Δx^2)
  as = -1/(Re*Δy^2)
  ae = -1/(Re*Δx^2)
  return ap,an,as,ae,aw
end
function BuildAdvectionDiffusionRHS!(Rp,ϕ,ψ,Δt,Δx,Δy,Ny,Nx,Re)
  # Time derivative
  Rp .= ϕ/Δt

  # Diffusion term (fully implicit)

  # Convection term
  for i in 2:Ny-1
    for j in 2:Nx-1
      ϕE = ϕ[i+0,j+1]; ϕW = ϕ[i+0,j-1]; ϕN = ϕ[i-1,j+0]; ϕS = ϕ[i+1,j+0]
      ψE = ψ[i+0,j+1]; ψW = ψ[i+0,j-1]; ψN = ψ[i-1,j+0]; ψS = ψ[i+1,j+0]

      u    = (ψN - ψS)/(2Δy); v    = -(ψE - ψW)/(2Δx)
      ∂ϕ∂y = (ϕN - ϕS)/(2Δy); ∂ϕ∂x = (ϕE - ϕW)/(2Δx)

      Rp[i,j] += - (u*∂ϕ∂x + v*∂ϕ∂y)
      # Rp[i,j] += (ψE - ψW)/(2Δx) * (ϕN - ϕS)/(2Δy) -
      #            (ψN - ψS)/(2Δy) * (ϕE - ϕW)/(2Δx)
    end
  end
end

function RecordHistory!(ϕ,ϕ_old,ϕ_hist)
  Δϕ = norm(ϕ - ϕ_old)
  ϕ_old .= ϕ
  push!(ϕ_hist,Δϕ)
  return(Δϕ)
end


# Linear mean function

@doc """
# Description
Constructor for the Linear mean function

m(x) = xᵀβ
# Arguments:
* `β::Vector{Float64}`: One coefficient for each dimension
""" ->
type MeanLin <: Mean
    β::Vector{Float64}
    dim::Int
    MeanLin(β::Vector{Float64}) = new(β, length(β))
end

mean(mLin::MeanLin,x::Matrix{Float64}) =  x'mLin.β

get_params(mLin::MeanLin) = mLin.β
get_param_names(::MeanLin) = [:β]
num_params(mLin::MeanLin) = mLin.dim

function set_params!(mLin::MeanLin, hyp::Vector{Float64})
    length(hyp) == mLin.dim || throw(ArgumentError("Linear mean function only has $(mLin.dim) parameters"))
    mLin.β = hyp
end

function grad_mean(mLin::MeanLin, x::Vector{Float64})
    dM_theta = x
    return dM_theta
end


module Day25

export day25

function powmod(a::T,n::T,m::T) where T<:Integer
    q = a
    product = T(1)
    while n > T(0)
        n,r = divrem(n,2)
        if r == T(1)
            product = mod(product * q, m)
        end
        q = mod(q^2, m)
    end
    return product
end

function invmod(a::T,m::T) where T<:Integer
    t = gcdx(a,m)
    if t[1] != T(1)
        return error("No inverse exists!")
    end
    return mod(t[2],m)
end

function dlog(g, y, p)
    q = Int(floor(sqrt(p)))
    table = Dict{Int,Int}([1 => 0])
    gq = powmod(g,q,p)
    current_power = gq
    for i in 1:div(p,q)
        table[current_power] = i
        current_power = mod(current_power*gq, p)
    end
    ginv = invmod(g,p)
    current_key = y
    for j in 0:q-1
        if haskey(table, current_key)
            return table[current_key] * q + j
        end
        current_key = mod(current_key*ginv, p)
    end
end

function day25()
    card = 10604480
    door = 4126658
    p = 20201227
    g = 7

    b = dlog(g, door, p)
    powmod(card,b,p)
end

end