using GraphPCA
using Base.Test

# X matrix
X1 = 10*randn(20,4)
X2 = X1 + randn(20,4)
X = [X1 X2]

#Construct the Laplacian matrix
L = eye(8)
for i in 1:4
  j = i + 4
  L[i,j] = -1
  L[j,i] = -1
end

λ = eigmax(Symmetric(X'*X))
η = eigmax(Symmetric(L))
eet = ones(8,8)
n = 8

#Pure PCA
XtX = Symmetric(I - X'*X/λ)
Dpca,Vpca = eig(XtX, 1:4)
Qpca = Vpca
Upca = X*Qpca
Xpca = Upca*Qpca'

#Pure Laplacian Embedding
Ln = Symmetric(L/η)
Dn, Vn = eig(Ln, 1:4)
Qlap = Vn
Ulap = X*Qlap
Xlap = Ulap*Qlap'

#Laplacian Embedding with Offset
# Lnp = Symmetric(L/η + eet/n)
# Dnp,Vnp = eig(Lnp, 1:4)
# Qlapp = Vnp
# Ulapp = X*Qlapp
# Xlapp = Ulapp*Qlapp'

GLPCA0 = GLPCA(X, 0, L, 4)
GLPCA1 = GLPCA(X, 1, L, 4)
@test GLPCA1.U*GLPCA1.Q' ≈ Xlap
@test GLPCA0.U*GLPCA0.Q' ≈ Xpca


#
# 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")


# Test for gaussjacobi().
tol = 1e-14
# Test a small n:
n = 42
a = -.1
b = .3
x, w = gaussjacobi(n, a, b)
@test all(length(x) == n) && all(length(w) == n)
@test abs(x[37] - 0.912883347814032) < tol
@test abs(w[37] - 0.046661910947553) < tol

@test_approx_eq dot(w,exp(x)) 2.7568520361985516

# Test a larger n (using ASY)
a = -.7
b = 1.3
n = 251
x, w = gaussjacobi(n, a, b)
@test all(length(x) == n) && all(length(w) == n)
@test abs(x[37] + 0.893103435898983) < tol
@test abs(w[37] - 1.962534523788093e-04) < tol

@test_approx_eq dot(w,exp(x)) 16.722957039404044

# Test n = 1:
a = 1.0
b = 2.0
x, w = gaussjacobi(1, a, b)
@test abs(x[1] - (b - a) / (a + b + 2)) < tol
@test abs(w[1] - 2^(a + b + 1) * beta(a + 1, b + 1)) < tol

@test_approx_eq dot(w,ones(x)) 1.3333333333333333

x, w = gaussjacobi(1013, .9, -.1)
@test abs(x[2] + 0.999986012231899) < tol
@test abs(x[13] + 0.999225722939832) < tol
@test abs(w[2] - 9.314674169892358e-05) < tol
@test abs(w[13] - 4.654651764553262e-04) < tol

@test_approx_eq dot(w,exp(x)) 1.6915068974063106

x, w = gaussjacobi(10013, .9, -.1)
@test abs(x[2]  +0.999999856605293) < tol
@test abs(x[13]  + 0.999992061552711) < tol
@test abs(w[2] - 1.509654630405615e-06) < tol
@test abs(w[13] - 7.548275262993863e-06) < tol

@test_approx_eq dot(w,exp(x)) 1.6915068974063106

# tests bug where row vectors were returned
@test isa(x,Vector{Float64})
@test isa(w,Vector{Float64})

# test last alpha and beta parameters:
x, w = gaussjacobi(100, 19., 21.)
@test abs(x[87] - 0.832211446176040) < tol
@test abs(w[50] - 0.064530500882703) < tol

# test for small alpha and beta:
x, w = gaussjacobi(10000, .1, .2)
@test abs(x[1] - -0.999999963363548) < tol
@test abs(w[500] - 2.183393039546711e-05) < tol

x,w=gaussjacobi(20,11.,0.)
# tests bug where row vectors were returned
@test isa(x,Vector{Float64})
@test isa(w,Vector{Float64})


# kron.jl
# test

using LinearMapsAA: LinearMapAA, LinearMapAM, LinearMapAO
using LinearAlgebra: I, Diagonal
using Test: @test, @testset

M1 = rand(ComplexF64, 3, 3)
M2 = rand(ComplexF64, 2, 2)

M12 = kron(M1, M2)
A1 = LinearMapAA(M1)
A2 = LinearMapAA(M2)
@testset "kron basics" begin
    @test kron(A1, A1) isa LinearMapAM
    for pair in ((A1,A2), (A1,M2), (M1,A2)) # all combinations
        AAk = kron(pair[1], pair[2])
        @test AAk isa LinearMapAM
        @test Matrix(AAk) == M12
        @test Matrix(AAk)' == Matrix(AAk')
    end
end

A3 = LinearMapAA(M1 ; odim=(1,size(M1,1)))
A4 = LinearMapAA(M2 ; idim=(size(M2,2),1))
@testset "kron(I,A)" begin
    K = kron(I(2), A4)
    M = kron(I(2), M2)
    @test K isa LinearMapAO
       @test Matrix(K) == M
       @test Matrix(K') == Matrix(K)'
end

@testset "kron(A,I)" begin
    K = kron(A4, I(2))
    M = kron(M2, I(2))
    @test K isa LinearMapAO
       @test Matrix(K) == M
       @test Matrix(K') == Matrix(K)'
end

@testset "kron(Ao,M)" begin
    @test kron(A3, M2) isa LinearMapAO
    @test kron(M1, A4) isa LinearMapAO
end

@testset "kron(Ao,Bo)" begin
    K = kron(A3, A4)
    M = kron(M1, M2)
    @test K isa LinearMapAO
       @test Matrix(K) == M
       @test Matrix(K') == Matrix(K)'
end

@testset "kron(A,D)" begin
    D = Diagonal(1:3)
    K = kron(A1, D)
    @test Matrix(K) == kron(M1, D)
    K = kron(D, A1)
    @test Matrix(K) == kron(D, M1)
end


### 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


using Test
import Calculus
import DSP
import NonlinearEquations
import Random
import SparseArrays

Random.seed!(0)

exorig = :(p[1] * x ^ 2 + p[2] * x + p[3])
ex, d = NonlinearEquations.replacerefswithsyms(exorig)
exold = NonlinearEquations.replacesymswithrefs(ex, d)
@test exorig == exold

ex = :(p[1] * x[1] ^ 2 + p[2] * x[2 * y + f(z)] + p[3])
diffs, refs = NonlinearEquations.differentiatewithrefs(ex, :x)
for i = 1:2
	if refs[i] == [:(2 * y + f(z))]
		@test Calculus.simplify(diffs[i]) == :(p[2])
	elseif refs[i] == [:(1)]
		@test Calculus.simplify(diffs[i]) == :(p[1] * (2 * x[1]))
	else
		error("the refs are bad")
	end
end

#formulate a quadratic equation
@NonlinearEquations.equations function quadeq(x, p; asdf=1)
	NonlinearEquations.setnumequations(1)
	NonlinearEquations.addterm(1, p[1] * x[1] ^ 2 + p[2] * x[1] + p[3])
end

#solve (x-1)(x+1)
p = [1.0, 0.0, -1.0]
root = NonlinearEquations.newton(x->quadeq_residuals(x, p), x->quadeq_x(x, p), [0.01])
@test root[1] ≈ 1.0
root = NonlinearEquations.newton(x->quadeq_residuals(x, p), x->quadeq_x(x, p), [100.0])
@test root[1] ≈ 1.0
root = NonlinearEquations.newton(x->quadeq_residuals(x, p), x->quadeq_x(x, p), [-0.01])
@test root[1] ≈ -1.0
root = NonlinearEquations.newton(x->quadeq_residuals(x, p), x->quadeq_x(x, p), [-23])
@test root[1] ≈ -1.0
#test the gradient
g(x, p) = x[1] * (p[1] + 2 * p[2] + 3 * p[3])
g_x(x, p) = [(p[1] + 2 * p[2] + 3 * p[3])]
g_p(x, p) = x[1] * [1, 2, 3]
f_x(x, p) = quadeq_x(x, p)
f_p(x, p) = quadeq_p(x, p)
g(p) = (-p[2] - sqrt(p[2] ^ 2 - 4 * p[1] * p[3])) / (2 * p[1]) * (p[1] + 2 * p[2] + 3 * p[3])
analytical_gradient = Calculus.gradient(g)
@test analytical_gradient(p) ≈ NonlinearEquations.gradient(root, p, g_x, g_p, f_x, f_p)

#solve (2x-3)(7x+5)
p = [14.0, -11.0, -15]
root = NonlinearEquations.newton(x->quadeq_residuals(x, p), x->quadeq_x(x, p), [1.0])
@test root[1] ≈ 3 / 2
root = NonlinearEquations.newton(x->quadeq_residuals(x, p), x->quadeq_x(x, p), [0.0])
@test root[1] ≈ -5 / 7

#formulate a quadratic equation in a different way
@NonlinearEquations.equations function quadeq2(x::T, p; asdf=1) where {T}
	NonlinearEquations.setnumequations(1)
	NonlinearEquations.addterm(1, p[1] * x[1] ^ 2)
	NonlinearEquations.addterm(1, p[2] * x[1])
	NonlinearEquations.addterm(1, p[3])
end

#solve (x-1)(x+1)
p = [1.0, 0.0, -1.0]
root = NonlinearEquations.newton(x->quadeq2_residuals(x, p), x->quadeq2_x(x, p), [0.01])
@test root[1] ≈ 1.0
root = NonlinearEquations.newton(x->quadeq2_residuals(x, p), x->quadeq2_x(x, p), [100.0])
@test root[1] ≈ 1.0
root = NonlinearEquations.newton(x->quadeq2_residuals(x, p), x->quadeq2_x(x, p), [-0.01])
@test root[1] ≈ -1.0
root = NonlinearEquations.newton(x->quadeq2_residuals(x, p), x->quadeq2_x(x, p), [-23])
@test root[1] ≈ -1.0

#solve (2x-3)(7x+5)
p = [14.0, -11.0, -15]
root = NonlinearEquations.newton(x->quadeq2_residuals(x, p), x->quadeq2_x(x, p), [0.0])
@test root[1] ≈ -5 / 7
root = NonlinearEquations.newton(x->quadeq2_residuals(x, p), x->quadeq2_x(x, p), [1.0])
@test root[1] ≈ 3 / 2

#formulate the steady-state diffusion equation with a heterogeneous diffusion coefficient
@NonlinearEquations.equations function diffusion(x, p, n)
	@assert length(x) == n
	@assert length(p) == n + 1
	NonlinearEquations.setnumequations(n)
	dx2 = 1 / (n + 1) ^ 2
	NonlinearEquations.addterm(1, ((p[1] + p[2]) * x[1] - p[2] * x[2] - p[1]) / dx2)
	for eqnum = 2:n - 1
		NonlinearEquations.addterm(eqnum, (-p[eqnum] * x[eqnum - 1] + (p[eqnum] + p[eqnum + 1]) * x[eqnum] - p[eqnum + 1] * x[eqnum + 1]) / dx2)
	end
	NonlinearEquations.addterm(n, ((p[end - 1] + p[end]) * x[end] - p[end - 1] * x[end - 1]) / dx2)
end

#solve the diffusion equation with head=1 at x=0 and head=0 at x=1 and a homogeneous diffusion coefficient
n = 10 ^ 3
p = ones(n + 1)
head = NonlinearEquations.newton(x->diffusion_residuals(x, p, n), x->diffusion_x(x, p, n), zeros(n); numiters=1)
dx2 = (1 / (n + 1)) ^ 2
otherhead = SparseArrays.spdiagm(-1=>-ones(n - 1) / dx2, 0=>2 * ones(n) / dx2, 1=>-ones(n - 1) / dx2) \ [[1 / dx2]; zeros(n - 1)]
@test head ≈ otherhead

#solve the diffusion equation with head=1 at x=0 and head=0 at x=1 and a heterogeneous diffusion coefficient
n = 10 ^ 3
windowsize = div(n, 10)
p = exp.(DSP.conv(0.1 * rand(n + windowsize + 1), ones(windowsize))[windowsize + 1:end - windowsize + 1])
head = NonlinearEquations.newton(x->diffusion_residuals(x, p, n), x->diffusion_x(x, p, n), zeros(n); numiters=1)
dx2 = (1 / (n + 1)) ^ 2
otherhead = SparseArrays.spdiagm(-1=>-p[2:n] / dx2, 0=>(p[1:n] + p[2:n + 1]) / dx2, 1=>-p[2:n] / dx2) \ [[p[1] / dx2]; zeros(n - 1)]
@test head ≈ otherhead

#formulate the boundary value problem u''(x)=u^2-sin^2(x)-sin(x) with u(0)=0 and u(1)=0, which has u(x)=sin(x) as a solution on the interval [0, π]
@NonlinearEquations.equations function nonlinearbvp(u, p, n)
	NonlinearEquations.setnumequations(n)
	dx2 = (pi / (n + 1)) ^ 2
	xs = range(0, pi; length=n + 2)[2:end - 1]
	NonlinearEquations.addterm(1, -(2 * u[1] - u[2] - 0) / dx2 - u[1] ^ 2 + sin(xs[1]) ^ 2 + sin(xs[1]))
	for eqnum = 2:n - 1
		NonlinearEquations.addterm(eqnum, -(2 * u[eqnum] - u[eqnum - 1] - u[eqnum + 1]) / dx2 - u[eqnum] ^ 2 + sin(xs[eqnum]) ^ 2 + sin(xs[eqnum]))
	end
	NonlinearEquations.addterm(n, -(2 * u[n] - u[n - 1] - 0 * sin(1)) / dx2 - u[n] ^ 2 + sin(xs[n]) ^ 2 + sin(xs[n]))
end
#now solve the boundary value problem
n = 10 ^ 4
p = []
u = NonlinearEquations.newton(x->nonlinearbvp_residuals(x, p, n), x->nonlinearbvp_u(x, p, n), zeros(n); numiters=100)
u_analytical = map(x->sin(x), collect(range(0, pi; length=n + 2))[2:end - 1])
@test u ≈ u_analytical


# 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 GreatCircle
using Geodesy: LatLon
using Base.Test

# One test decimal degree is 111000m
latitude = 40.0
longitude = -76.0
position = LatLon(lat=latitude, lon=longitude)

"""
    Great Circle Tests
"""
azimuth = 90.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone to the right
@test new_gc["longitude"] > longitude + 0.9

a = great_circle(111320., azimuth, 0., 0.)
@test abs(a["latitude"]-0.)<=eps(typeof(a["latitude"]))
for i = 1:2
    a = great_circle(111320., azimuth, a["latitude"], a["longitude"])
    @test abs(a["latitude"]-0.)<=eps(typeof(a["latitude"]))
end
aposition = LatLon(0., 0.)
a = great_circle(111320., azimuth, aposition)
@test abs(a.lat-0.)<=eps(typeof(a.lat))
for i = 1:2
    a = great_circle(111320., azimuth, a)
    @test abs(a.lat-0.)<=eps(typeof(a.lat))
end

azimuth = 270.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone to the left
@test new_gc["longitude"] < longitude - 0.9

azimuth = 180.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone down
@test new_gc["latitude"] < latitude - 0.9

azimuth = 0.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone up
@test new_gc["latitude"] > latitude + 0.9

azimuth = 315.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone up and to the left
@test new_gc["latitude"] > latitude + 0.45
@test new_gc["longitude"] < longitude - 0.45

new_gc = great_circle(111000., azimuth, position)
@test new_gc.lat > latitude + 0.45
@test new_gc.lon < longitude - 0.45

# One decimal degree is 111000m
latitude = collect(40.0:5.:60.0)
longitude = collect(-40.0:-5.:-60.0)

azimuth = ones(longitude) * 90.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone to the right
@test all(new_gc["longitude"] .> longitude + 0.9)

azimuth = 270.
distance = 111000. * ones(longitude)
new_gc = great_circle(distance, azimuth, latitude, longitude)
# We should have gone to the left
@test all(new_gc["longitude"] .< longitude - 0.9)

azimuth = 180.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone down
@test all(new_gc["latitude"] .< latitude - 0.9)

azimuth = 0.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone up
@test all(new_gc["latitude"] .> latitude + 0.9)

azimuth = 315.
new_gc = great_circle(111000., azimuth, latitude, longitude)
# We should have gone up and to the left
@test all(new_gc["latitude"] .> latitude + 0.45)
@test all(new_gc["longitude"] .< longitude - 0.45)

"""
    Great Distance Tests
"""
# One decimal degree at the equator is about 111.32km
latitude_start  = 0.
latitude_end    = 0.
longitude_start = 50.
longitude_end   = 52.
gd = great_distance(latitude_start, longitude_start, latitude_end, longitude_end)
@test round(gd["distance"] ./ 1000, 2) .== 111.32 .* 2

a = LatLon(lat=latitude_start, lon=longitude_start)
b = LatLon(lat=latitude_end, lon=longitude_end)
gd = great_distance(a, b)
@test round(gd["distance"] ./ 1000, 2) .== 111.32 .* 2

# One decimal degree is 111000m
latitude_start  = 0.
latitude_end    = 0.
longitude_start = [49., 75.]
longitude_end   = [50., 76.]
gd = great_distance(latitude_start, longitude_start, latitude_end, longitude_end)
@test all( round(gd["distance"] ./ 1000, 2) .== 111.32 )


using DiffEqOperators
# using Plots

n = 10000
@show n
x = range(0.0; length = n, stop = 2π)
dx = x[2]-x[1]
y = exp.(π*x)
y_ = y[2:(end-1)]

for dor in 1:4, aor in 2:6

    D1 = CenteredDifference(dor,aor,dx,length(x))

    #test result
    @show dor
    @show aor
    #take derivatives
    err_abs = abs.(D1*y .- (π^dor)*y_)
    err_percent = 100*err_abs./abs.(((π^dor)*y_))
    max_error = maximum(err_percent) # test operator with known derivative of exp(kx)
    avg_error = sum(err_percent)/length(err_percent)

    @show max_error
    @show avg_error
    #plot(x[2:(end-1)], err_percent, title="Percentage Error, n=$n aor = $aor, dor = $dor")

    #test result
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


using ControlSystems;

A=[0 1;1 0];
B=[0;1];
Q=[2 2;2 2];
R=8;

l=lqr(A,B,Q,R)

care(A,B,Q,R)


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


export PRISMA, PrismaParams, PrismaStepsize

function PRISMA(x, # starting point
				Lf, # Lipshitz constant of f
				grad_f,
				prox_g,
				prox_h,
				obj,
				params::OptParams=PrismaParams(PrismaStepsize(), 10))

	# initialize
	ssr      = params.stepsizerule
	if ssr.initial_stepsize == Inf
		R = sqrt(obj(x)) # estimate of distance to solution
		ssr.initial_stepsize = 2*R/Lf
	end
	y        = copy(x)
	xkpp     = copy(x)
	betakpp  = step(ssr)
	Lkpp     = Lf + 1/betakpp
	thetak   = 1
	thetakpp = 1 # XXX see orabona's code

	if params.verbose > 0
		@printf("%10s%12s\n", "iter", "obj")
		@printf("%10d%12.4e\n", 0, obj(xkpp))
	end

	# iterate
	for k=1:params.maxiter
		betak, betakpp    = betakpp, step(ssr)
		Lk, Lkpp          = Lkpp, Lf + 1/betakpp
		thetak, thetakpp  = thetakpp, 2/(1+sqrt(1+4Lkpp/thetak^2/Lk))
		xk, xkpp          = xkpp, prox_h((1-1/Lk/betak)*y - grad_f(y)/Lk + prox_g(y, betak)/Lk/betak, 1/Lk)
		y                 = xkpp + thetakpp*(1/thetak - 1)*(xkpp - xk)
		if params.verbose > 0
			@printf("%10d%12.4e\n", k, obj(xkpp))
		end
		if stop(params, xk, xkpp)
			break
		end
	end

	return xkpp
end

type PrismaStepsize<:StepSizeRule
	initial_stepsize::AbstractFloat
	iteration_counter::Int
end
function step(s::PrismaStepsize, args...; kwargs...)
	s.iteration_counter += 1
	return s.initial_stepsize/(s.iteration_counter)
end
PrismaStepsize() = PrismaStepsize(Inf)
PrismaStepsize(initial_stepsize) = PrismaStepsize(initial_stepsize, 0)

type PrismaParams<:OptParams
	stepsizerule::StepSizeRule
	maxiter::Int
	verbose::Int
	reltol::Float64
end
PrismaParams(;stepsize::StepSizeRule=PrismaStepsize(),
	          maxiter::Int=100,
	          verbose::Int=1,
	          reltol::Float64=1e-5) = PrismaParams(stepsize,maxiter,verbose,reltol)
stop(s::PrismaParams, xk, xkpp) = vecnorm(xk - xkpp)/vecnorm(xk) < s.reltol ? true : false


#=
We check the deposition step (i.e. computing \rho) works as expected.
=#
using Test

@testset "PIC deposition : uniform    " begin
    using LinearAlgebra
    using VlasovSolvers
    # using Plots

    include("PIC.jl")

    function samples_verif1(nsamples, kx::Float64)
        x0 = Array{Float64}(undef, nsamples)
        seq = SobolSeq(1)
        for i = 1:nsamples
            r = Sobol.next!(seq)
            x0[i] = 2π/kx * r[1]
        end
        v0 = rand(Uniform(-10, 10), nsamples)
        wei = 2π/kx/nsamples .* ones(Float64, nsamples)
        return x0, v0, wei
    end

    kx = 0.5
    L = 2π / kx
    np = 100000 # nb de particules
    nx = 32   # nb de mailles pour le calcul de l energie
    mesh1 = OneDGrid(CPU(), nx, 0, L)

    (x0, v0, w0) = samples_verif1(np, kx)
    p = Particles(x0, v0, w0, np)

    ρ, ρ_tot = compute_rho(p, mesh1)

    # arrx = mesh1.start:mesh1.step:(mesh1.stop - mesh1.step)
    # plot(arrx, ρ)

    @test maximum(abs.(ρ .- 1.)) < 10^-3
    @test ρ_tot ≈ 1.0
end

@testset "PIC deposition : 1 + 0.99cos" begin
    using LinearAlgebra, Roots
    using VlasovSolvers
    # using Plots

    include("PIC.jl")

    function samples_verif2(nsamples)
        x0 = Array{Float64}(undef, nsamples)
        seq = SobolSeq(1)
        for i = 1:nsamples
            r = Sobol.next!(seq)
            resx = find_zero(x-> x + 0.99 / kx * sin(kx*x) - r[1]*2π/kx, 0.0)
            x0[i] = resx
        end
        v0 = rand(Uniform(-10, 10), nsamples)
        wei = 2π/kx/nsamples .* ones(Float64, nsamples)
        return x0, v0, wei
    end

    kx = 0.5
    L = 2π / kx
    np = 100000 # nb de particules
    nx = 128   # nb de mailles pour le calcul de l energie
    mesh1 = OneDGrid(CPU(), nx, 0, L)

    (x0, v0, w0) = samples_verif2(np)
    p = Particles(x0, v0, w0, np)

    ρ, ρ_tot = compute_rho(p, mesh1)

    arrx = mesh1.start:mesh1.step:(mesh1.stop - mesh1.step)

    @test maximum(abs.(ρ - (1 .+ 0.99cos.(kx.*arrx)))) < 5.10^-3
    @test ρ_tot ≈ 1.0
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


# Based on an example from
#	https://jump.dev/JuMP.jl/previews/PR2476/examples/Semidefinite%20programs/max_cut_sdp/
# which in turn uses methods from
#	Goemans, M. X., & Williamson, D. P. (1995). Improved approximation
#	algorithms for maximum cut and satisfiability problems using semidefinite
#	programming. Journal of the ACM (JACM), 42(6), 1115-1145.

using SparseArrays
using LinearAlgebra

# Model the max-cut SDP relaxation --------------------------------------------

cut_value(C, cut) =	sum(C .* (cut * cut'))
laplacian(W) = spdiagm(vec(sum(W, dims=1))) - W

function extract_cut(X)
	rnd = normalize(rand(size(X, 1)))

	factorization = cholesky(Hermitian(X), Val(true); check=false)
	V = (factorization.P * factorization.L)'
	vs = map(normalize, eachcol(V))

	map(x -> dot(rnd, x) |> sign, vs)
end

function max_cut_sdp_model(weights)
	n = size(weights, 1)

	C = laplacian(weights) / 4
	As = [sparse([i], [i], [1], n, n) for i in 1:n]
	b = ones(n) / n

	C, As, b
end

function solve_maxcut(W)
	n = size(W, 1)

	C, As, b = max_cut_sdp_model(W)
	X = sketchy_cgal(normalize(-C), As, b; R=n, info_io=devnull)

	cut = extract_cut(X)
	value = cut_value(C, cut)

	cut, value
end

# Examples --------------------------------------------------------------------

function maxcut_example_1()
	W = [
		0 5;
		5 0;
	]

	cut, value = solve_maxcut(W)

	@test isapprox(value, 5)
	@test cut[1] != cut[2]
end

function maxcut_example_2()
	W = [
		0 5 7 6;
		5 0 0 1;
		7 0 0 1;
		6 1 1 0;
	]

	cut, value = solve_maxcut(W)

	@test isapprox(value, 18)
	@test cut[1] != cut[2]
	@test cut[2] == cut[3] == cut[4]
end

function maxcut_example_3()
	W = [
		0 1 5 0;
		1 0 0 9;
		5 0 0 2;
		0 9 2 0;
	]

	cut, value = solve_maxcut(W)

	@test isapprox(value, 17)
	@test cut[1] == cut[4]
	@test cut[2] == cut[3]
	@test cut[1] != cut[2]
end


__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


export yin

"""
    yin(sig::Vector, sr::Int; w_len::Int = 512, f_step::Int = 256, f0_min = 100, f0_max = 500,
        harmonic_threshold = 0.1,
        difference_function::Function = difference_function_original, kwargs...) -> F0s, frame_times

Estimates the fundamental frequency (F0) of the signal `sig` using the YIN algorithm [^1].
The signal `sig` is a vector of points uniformly sampled at a rate `sr`.

## Keyword arguments
* `w_len`: size of the analysis window [samples == number of points]
* `f_step`: size of the lag between two consecutive frames [samples == number of points]
* `f0_min`: Minimum fundamental frequency that can be detected [linear frequency]
* `f0_max`: Maximum fundamental frequency that can be detected [linear frequency]
* `harmonic_threshold`: Threshold of detection. The algorithm returns the first minimum of
  the CMNDF function below this threshold.
* `diffference_function`: The difference function to be used (by default
  [`differenceFunction_original`](@ref)).

## Description
The YIN algorithm [^CheveigneYIN2002] estimates the signal's fundamental frequency `F0` by basically
looking for the period `τ0`  which minimizes the signal's autocorrelation. This
autocorrelation is calculated for signal segments (frames), composed of two windows of
length `w_len`. Each window is separated by a distance `τ`, and the idea is that the
distance which minimizes the pairwise difference between each window is considered to be the
fundamental period `τ0` of that frame.

More precisely, the algorithm first computes the cumulative mean normalized difference
function (MNDF) between two windows of a frame for several candidate periods `τ` ranging
from `τ_min=sr/f0_max` to `τ_max=sr/f0_min`. The MNDF is defined as
```math
d_t^\\prime(\\tau) = \\begin{cases}
        1 & \\text{if} ~ \\tau=0 \\\\
        d_t(\\tau)/\\left[{(\\frac 1 \\tau) \\sum_{j=1}^{\\tau} d_{t}(j)}\\right] & \\text{otherwise}
        \\end{cases}
```
where `d_t` is the difference function:
```math
d_t(\\tau) = \\sum_{j=1}^W (x_j - x_{j+\\tau})^2
```

It then refines the local minima of the MNDF using parabolic (quadratic) interpolation. This
is done by taking each minima, along with their first neighbor points, and finding the
minimum of the corresponding interpolated parabola. The MNDF minima are substituted by the
interpolation minima. Finally, the algorithm chooses the minimum with the smallest period
and with a corresponding MNDF below the `harmonic threshold`. If this doesn't exist, it
chooses the period corresponding to the global minimum. It repeats this for frames starting
at the first signal point, and separated by a distance `f_step` (frames can overlap), and
returns the vector of frequencies `F0=sr/τ0` for each frame, along with the start times of
each frame.

As a note, the physical unit of the frequency is 1/[time], where [time] is decided by the
sampling rate `sr`. If, for instance, the sampling rate is over seconds, then the frequency
is in Hertz.

[^CheveigneYIN2002]: De Cheveigné, A., & Kawahara, H. (2002). YIN, a fundamental frequency estimator for
speech and music. The Journal of the Acoustical Society of America, 111(4), 1917-1930.
"""
function yin(sig::Vector, sr::Int; w_len::Int = 512, f_step::Int = 256, f0_min = 100,
    f0_max = 500, harmonic_threshold = 0.1,
    difference_function::Function = difference_function_original, kwargs...)

    τ_min = floor(Int64, sr / f0_max)
    τ_max = floor(Int64, sr / f0_min)

    frame_times = range(1, length(sig) - w_len - τ_max, step=f_step)  # time values for start of  each analysis window
    # times = frame_times ./ eltype(sig)(sr)

    F0s = zeros(Float64, length(frame_times))

    for (i, t) in enumerate(frame_times)
        frame = sig[ (t) : (t + τ_max+ w_len-1) ]
        df = difference_function(frame, w_len, τ_max)
        cmdf = cumulative_mean_normalized_difference_function(df)
        y_refined, τ_refined, τ_indices = refine_local_minima(cmdf)
        idx_localminimum = absolute_threshold(y_refined, τ_min, τ_max, harmonic_threshold)
        τ0 = τ_refined[idx_localminimum]
        F0s[i] = sr / τ0

    end
    return  F0s, frame_times
end

"""
    difference_function_original(x, W, τmax) -> df
Computes the difference function of `x`. `W` is the window size, and `τmax` is the maximum
    period. This corresponds to equation (6) in [^1]:
```math
d_t(\\tau) = \\sum_{j=1}^W (x_j - x_{j+\\tau})^2
```
"""
function difference_function_original(x, W, τmax)
    df = zeros(eltype(x), τmax+1) #df corresponds to τ values from 0 to τ_max
    for τ in 1:τmax
        for j in 1:W
            df[τ+1] += (x[j] - x[j + τ]) ^ 2
        end
    end
    return df
end

"""
Compute cumulative mean normalized difference function (CMND), starting from the difference
    function `df`. This corresponds to equation (8) in [1]
"""
function cumulative_mean_normalized_difference_function(df)
    N = length(df)
    cmndf = df[2:end] .* range(1, N-1, step=1) ./ Float64.(cumsum(df[2:end]))
    return [1.0; cmndf]
end

"""
Returns the refined local minima of `y`, along with their refined `x` value and the
corresponding indices.  For each minimum, by the minimum of the parabola obtained by
interpolated the minimum with its first neighbors. Also returns the indices of `y`
(`x_nominal`) and the value `x` corresponding to each minima.
"""
function refine_local_minima(y)
    x_nominal = 0:length(y)-1 #nominal τ values: 0 to τ_max (cf. difference_function)
    idxs_local_minima = local_minima(y)
    if isempty(idxs_local_minima)
         @warn "No local minima found for the cumulative mean difference function. Adjusting
         the values of the minimum and maximum frequencies may fix this."
    end
    xv, yv = parabolic_interpolation(y, idxs_local_minima) #xv,yv are correction for the local minima
    y_locmin_refined = yv
    x_locmin_real= x_nominal[idxs_local_minima] .+ xv
    x_locmin_nominal = x_nominal[idxs_local_minima]
    return y_locmin_refined, x_locmin_real, x_locmin_nominal
end

"""
Applies the threshold step described in [^1]. It returns index (period) corresponding to the
first minimum below the `harmonic_threshod`. If that doesn't exist, it regurns the index of the
global minimum.
"""
function absolute_threshold(localminima, τ_min, τ_max, harmonic_threshold=0.1)
    for (idx, localminimum) in enumerate(localminima)
        if localminimum ≤ harmonic_threshold
            return idx
        end
    end
    return argmin(localminima)
end

"""
Calculates the parabolic (quadratic) interpolation for all the indixes of `y` given in
idxs_interpolate, and returns the coordinates of the respective parabola's minimum.
Assumes space all adjacent x values of `y` is 1.
"""
function parabolic_interpolation(y, idxs_interpolate)
    #separate x into triplets [x1,x2,x3].
    x1 = @view y[idxs_interpolate .- 1]
    x2 = @view y[idxs_interpolate]
    x3 = @view y[idxs_interpolate .+ 1]
    #calculate the vertix coordinates (xv, yv) for each triplet
    a = @. (x1 - 2x2 + x3)/2
    b = @. (x3 - x1) / 2
    xv = @. -b/2a
    yv = @. x2 - (b^2 /4a)
    return xv, yv
end

"""
Quickly written version for finding local minima by comparing first neighbors only.
#TODO: More detailed implementations available in Peaks.jl or Iamges.jl, maybe this should
be replaced by them.
"""
function local_minima(x)
    x1 = @view x[1:end - 2]
    x2 = @view x[2:end - 1]
    x3 = @view x[3:end]
    collect(1:length(x2))[x1 .> x2 .< x3] .+ 1
end


@testset "Covariances" begin
  comps = rand(Composition{5}, 100)
  data  = reduce(hcat, components.(comps))
  table = (; zip([:N₂,:O₂,:CO₂,:S,:C], eachrow(data))...)

  Τ = variation(table)
  Σ = alrcov(table)
  Γ = clrcov(table)
  A = lrarray(table)

  # Testing matrix forms transformations (Aitchison 1986, pages 82 and 83)
  ## Test matrix transformation from Τ to Σ (4.25)
  @test Σ ≈ -0.5 * F * Τ * F'

  ## Test matrix transformation from Τ to Γ (4.26)
  @test Γ ≈ -0.5 * G * Τ * G

  ## Test matrix transformation from Σ to Τ (4.27)
  Γ₀ = F' * (inv(H) * Σ * inv(H)) * F
  @test Τ ≈ J * Diagonal(Γ₀) + Diagonal(Γ₀) * J - 2 * Γ₀

  ## Test matrix transformation from Σ to Γ (4.28)
  @test Γ ≈ F' * (inv(H) * Σ * inv(H)) * F

  ## Test matrix transformation from Γ to Τ (4.29)
  @test Τ ≈ J * Diagonal(Γ) + Diagonal(Γ) * J - 2 * Γ

  ## Test matrix transformation from Γ to Σ (4.30)
  @test Σ ≈ F * Γ * F'
end


module SphericalHarmonics_Mod
    using LegendrePolynomials

    export Calc_SphericalHarmonics

    function Calc_SphericalHarmonics(l::Int64,m::Int64,n::Int64 = 50)

        #Create meshgrid
        theta = range(0,pi,length=n)' .* ones(n,1)
        phi = range(0,2*pi,length=n) .* ones(1,n)

        #Calculate spherical harmonics
        Y = harmonicY(l,m,theta,phi,"real")

        # Transform from spherical to cartesian coordinates
        x,y,z = sph2cart(phi, pi/2 .- theta, abs.(Y))

        return x,y,z,Y
    end

    function harmonicY(l::Int64,m::Int64,th::Matrix{Float64},phi::Matrix{Float64},type::String)
        # Calculate the spherical harmonic Ylm
        # Translation of: Javier Montalt Tordera (2022). Spherical Harmonics (https://github.com/jmontalt/harmonicY/releases/tag/v2.0.1), GitHub. Retrieved January 19, 2022.

        @assert abs(m) <= l "abs(m) > l!"

        isoddm = isodd(m)
        isnegm = m < 0

        # Quantum mechanical norm, se Wiki.
        # Factorials can overflow Float64 for l > 19, hence "big".
        C = (-1)^m*sqrt((2*l+1)/(4*pi)*factorial(big(l-m))/factorial(big(l+m)))

        C = Float64(C)

        m = abs(m)
        P = Plm.(cos.(th),l,m)

        if type=="real"
            if isnegm
                E = sin.(m*phi)
            else
                E = cos.(m*phi)
            end
        else
            E = exp.(1im*m*phi)
            if isnegm
                conj!(E)
                if isoddm
                    E = -E
                end
            end
        end

        # Surface spherical harmonics
        Y = C * P .* E

        return Y
    end

    function sph2cart(az::Matrix{Float64},el::Matrix{Float64},r::Matrix{Float64})
        #Spherical to cartesian

        z = r .* sin.(el)
        rcoselev = r .* cos.(el)
        x = rcoselev .* cos.(az)
        y = rcoselev .* sin.(az)

        return x,y,z
    end
end


using ModelingToolkit
using OrdinaryDiffEq

"""
solves STG calcium model (Liu 1998)

The log(3000.0/Ca) in the algorithm doesn't play nicely with Zygote. Don't know why. I've maxed it with zero to make it work.

All neurons have following channels (with some conductances at 0:)
NaV
CaT
CaS
H
Kd
Leak
"""

function  CalciumNeuron(input)

    @parameters t
    D = Differential(t)

    @parameters eNa eh eK eleak tauCa gNabar gCaSbar gCaTbar gKabar gKCabar gKdrbar ghbar gleak

    paramvars = [eNa, eh, eK, eleak, tauCa, gNabar, gCaSbar, gCaTbar, gKabar, gKCabar, gKdrbar, ghbar, gleak]

    @variables V(t) Ca(t) mNa(t) hNa(t) mCaS(t) hCaS(t) mCaT(t) hCaT(t) mIh(t) mKa(t) hKa(t) mKCa(t) mKdr(t)
    statevars = [V, Ca, mNa, hNa, mCaS, hCaS, mCaT, hCaT, mIh, mKa, hKa, mKCa, mKdr]

    # calcium reversal potential
    # eCa = (500.0)*(8.6174e-5)*(283.15)*(log(max((3000.0/Ca),1e-5)))
    eCa = (500.0)*(8.6174e-5)*(283.15)*(log((3000.0/Ca)))

    # currents
    INa=  gNabar*mNa^3*hNa*(eNa - V)
    ICaS= gCaSbar*mCaS^3*hCaS*(eCa - V)
    ICaT= gCaTbar*mCaT^3*hCaT*(eCa - V)
    Ih= ghbar*mIh*(eh - V)
    IKa= gKabar*mKa^3*hKa*(eK - V)
    IKCa=  gKCabar*mKCa^4*(eK - V)
    IKdr=  gKdrbar*mKdr^4*(eK - V)
    Ileak= gleak*(eleak - V)

    # calcium
    Ca_inf = 0.05 + 0.94*(ICaS + ICaT);

    # gating
    Na_m_inf=  1.0./(1.0+exp((V+25.5)./-5.29));
    Na_h_inf=  1.0./(1.0+exp((V+48.9)./5.18));
    Na_tau_m=  1.32 - 1.26./(1+exp((V+120.0)./-25.0));
    Na_tau_h=  (0.67./(1.0+exp((V+62.9)./-10.0))).*(1.5+1.0./(1.0+exp((V+34.9)./3.6)));

    CaS_m_inf=  1.0./(1.0+exp((V+33.0)./-8.1));
    CaS_h_inf=  1.0./(1.0+exp((V+60.0)./6.2));
    CaS_tau_m=  1.4 + 7.0./(exp((V+27.0)./10.0) + exp((V+70.0)./-13.0));
    CaS_tau_h=  60.0 + 150.0./(exp((V+55.0)./9.0) + exp((V+65.0)./-16.0));

    CaT_m_inf=  1.0./(1.0 + exp((V+27.1)./-7.2));
    CaT_h_inf=  1.0./(1.0 + exp((V+32.1)./5.5));
    CaT_tau_m=  21.7 - 21.3./(1.0 + exp((V+68.1)./-20.5));
    CaT_tau_h=  105.0 - 89.8./(1.0 + exp((V+55.0)./-16.9));

    Ih_m_inf=  1.0./(1.0+exp((V+70.0)./6.0));
    Ih_tau_m=  (272.0 + 1499.0./(1.0+exp((V+42.2)./-8.73)));

    Ka_m_inf=  1.0./(1.0+exp((V+27.2)./-8.7));
    Ka_h_inf=  1.0./(1.0+exp((V+56.9)./4.9));
    Ka_tau_m=  11.6 - 10.4./(1.0+exp((V+32.9)./-15.2));
    Ka_tau_h=  38.6 - 29.2./(1.0+exp((V+38.9)./-26.5));

    KCa_m_inf=  (Ca./(Ca+3.0))./(1.0+exp((V+28.3)./-12.6));
    KCa_tau_m=  90.3 - 75.1./(1.0+exp((V+46.0)./-22.7));

    Kdr_m_inf=  1.0./(1.0+exp((V+12.3)./-11.8));
    Kdr_tau_m=  7.2 - 6.4./(1.0+exp((V+28.3)./-19.2));

    eqs = [
    D(V) ~              INa + ICaS + ICaT + Ih + IKa + IKCa + IKdr + Ileak,
    D(Ca) ~                   (1/tauCa)*(Ca_inf - Ca),
    D(mNa) ~                   (1/Na_tau_m)*(Na_m_inf - mNa),
    D(hNa) ~                   (1/Na_tau_h)*(Na_h_inf - hNa),
    D(mCaS) ~                   (1/CaS_tau_m)*(CaS_m_inf - mCaS),
    D(hCaS) ~                   (1/CaS_tau_h)*(CaS_h_inf - hCaS),
    D(mCaT) ~                   (1/CaT_tau_m)*(CaT_m_inf - mCaT),
    D(hCaT) ~                   (1/CaT_tau_h)*(CaT_h_inf - hCaT),
    D(mIh) ~                   (1/Ih_tau_m)*(Ih_m_inf - mIh),
    D(mKa) ~                  (1/Ka_tau_m)*(Ka_m_inf - mKa),
    D(hKa) ~                   (1/Ka_tau_h)*(Ka_h_inf - hKa),
    D(mKCa) ~                  (1/KCa_tau_m)*(KCa_m_inf - mKCa),
    D(mKdr) ~                   (1/Kdr_tau_m)*(Kdr_m_inf - mKdr)
    ]

    ps = paramvars .=> [50,-20,-80,-50,20,100,3,1.3,5,10,20,0.5,0.01]
    ics = statevars .=> [-60,0.05,0,0,0,0,0,0,0,0,0,0,0]

    od = ODESystem(eqs, t, statevars, paramvars)
    tspan = (0.,500.)
    return od, ics, tspan, ps
end

# function predict(p)
#     return Array(solve(prob, Tsit5(), p=p, saveat=tsteps, sensitivity=BacksolveAdjoint(;autojacvec=true)))
# end

# function loss(p)
#     prediction = predict(p)
#     return sum(abs2, prediction - data)
# end


"""
direct product of two or more discrete functions

if we have:
    ``f_1 : D_1 → R_1``
    `` …``
    ``f_k : D_k → R_k``
then FunProduct((f1, …, fk)) is F
    F : D1 × … × Dk → R1 × … × Rk
"""
struct FunProduct{N, T, D, R} <: CompositeFunction
    factors :: NTuple{N, T}
    dom :: D
    rng :: R
end

const FProd = FunProduct

domain(fp :: FProd) = fp.dom
codomain(fp :: FProd) = fp.rng
factors(fp :: FProd) = fp.factors

# creating functional product from dom and codom
function FProd(dom :: AbstractProduct, codom :: AbstractProduct)
    fs = Tuple(ResFun(dom[i], codom[i]) for i in 1 : ndims(dom))
    return FProd(fs, dom, codom)
end

# creating functional product from two tuples
# of equal length!

function FProd(from :: NTuple{N, Int}, to :: NTuple{N, Int}) where N
    fromto = zip(from, to)
    funs = tuple(map(ResidueFunction, fromto)...)
    return FProd(funs)
end

function FProd(factors)
    dom = DirectProduct(map(domain, factors))
    rng = DirectProduct(map(codomain, factors))
    return FProd(factors, dom, rng)
end

FProd(factors...) = FProd(tuple(factors...))

function (fp :: FProd{N, T, D, R})(v) where {N,T,D,R}
    funs = factors(fp)
    n = length(v)
    return NTuple{N}(funs[i](v[i]) for i in 1 : n)
end

function Base.setindex!(fp :: FProd, v :: NTuple, k :: NTuple)
    fs = factors(fp)
    for i in 1 : size(fp)
        fs[i][k[i]] = v[i]
    end
end
(fs :: FProd)(v...) = fs(tuple(v...))