### A Pluto.jl notebook ###
# v0.9.11

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

# ╔═╡ 77489ff6-bd82-11ea-2d09-8737634cc579
using Random

# ╔═╡ e9b93d52-bd82-11ea-15c9-699eeebaa243
using Plots

# ╔═╡ 7eab739e-bdba-11ea-3ff3-fd88b08946d3
using LaTeXStrings

# ╔═╡ 508fed7e-bdbe-11ea-1d75-8935a47d06d8
using Statistics

# ╔═╡ fc08751c-bd79-11ea-072b-33a239fa0605
md"# Stochastic Thinking"

# ╔═╡ f895f5da-bd7c-11ea-2250-efdcd9d49a38
md"## Modelling the recovery process using probability
Recall code: $I_{n+1} = (1 + c) I_n = \lambda * I_n$"

# ╔═╡ 65f53264-bd7d-11ea-03cc-9513654c222c
function run_infection(I₀, λ, T=20)
	I = zeros(T)
	I[1] = I₀

	for n ∈ 1:T-1
		I[n+1] = λ * I[n]
	end

	return I
end

# ╔═╡ 2dce78b0-bd7e-11ea-0381-d16883397ebd
run_infection(1.0, 0.1)

# ╔═╡ 38ad903e-bd7e-11ea-08b6-f9a9a086260f
md"## Adding new data to extend a vector"

# ╔═╡ c7c0a170-bd7f-11ea-3b68-a9ae0c8ccfde
v = [1.0]

# ╔═╡ e40b10e0-bd7f-11ea-2883-bf485b88d6ad
push!(v, 7.0) # ! means the function modifies its first argument

# ╔═╡ b9db886e-bd80-11ea-1347-5d5415a22dd5
v

# ╔═╡ c427a4fc-bd80-11ea-365a-59fd43802b47
[v; 10]

# ╔═╡ 92666d14-bd81-11ea-2ba7-e181ffa9f476
pushfirst!(v, 8) # deque - double endedn queue

# ╔═╡ 0f3c4b1c-bd80-11ea-0fd0-df1c82c3ed1c
function run_infection_vectorized(I₀, λ, T=20)
	Is = [I₀]
	I = I₀

	for n ∈ 1:T-1
		I_next = λ * I

		push!(Is, I_next)

		I = I_next
	end

	return Is
end

# ╔═╡ 71759792-bd81-11ea-1cbe-35ca8456b6a1
run_infection_vectorized(1.0, 1.1)

# ╔═╡ 7a610d46-bd81-11ea-0ff9-5b97563bce1e
md"## Randomness"

# ╔═╡ 29e900e6-bd82-11ea-0344-2b2276b99422
rand()

# ╔═╡ 8cec500a-bd82-11ea-2fbb-03dec9457d7e
Random.seed!(3);

# ╔═╡ a1c76a96-bd82-11ea-22a0-d9110840ecf3
rand()

# ╔═╡ d39e0c64-bd82-11ea-188d-6f762d4dd688
rand(10)

# ╔═╡ ed82a2b8-bd82-11ea-1d5c-8759b29ce1cc
scatter(rand(50))

# ╔═╡ 0682ca0c-bd83-11ea-0afa-8794eb7bd9b8
begin
	local r = rand(50)

	scatter(r, 0.5 * one.(r), ylim=(0, 1))
end

# ╔═╡ be2d6b30-bd83-11ea-2683-49e64cbccb3c
@bind num_points html"<input type='range' min=1 max=100>"

# ╔═╡ 811488a0-bdab-11ea-0f9d-b1bc9f854e07
#TODO coloring exercise
begin
	local r = rand(num_points)

	scatter(r[1:num_points],0.5 .* one.(r[1:num_points]),ylim=(0, 1),xlim=(0, 1))
end

# ╔═╡ 89984d06-bdab-11ea-3ff4-55786243da18
md"## Event with probability $p$"

# ╔═╡ a8196c14-bdac-11ea-3b3d-f744310c289d
begin
	local p = 0.25
	local r = rand()

	if r < p
		true
	else
		false
	end
end

# ╔═╡ cfac1f86-bdac-11ea-316c-a9abea7ed509
md"Bernoulli Trial"

# ╔═╡ 46de61a4-bdad-11ea-3aaf-a5f9b96e533f
function bernoulli(p)
	r = rand()

	return r < p
end

# ╔═╡ 8a20f6b6-bdad-11ea-0ef5-af788315bd82
2 < 3 ? true : false # ternary operator

# ╔═╡ bc82b248-bdad-11ea-3ad5-019463cf2ef1
2 < 3

# ╔═╡ c7552f0e-bdad-11ea-3fe8-fd3cf143ff30
begin
	local p = 0.25

	trials = [bernoulli(p) for i ∈ 1:100]
end

# ╔═╡ fe9569d2-bdad-11ea-0f72-09084f78f692
scatter(trials)

# ╔═╡ 38c6f580-bdae-11ea-2cbe-935e598ee04c
trials

# ╔═╡ 3a406aea-bdae-11ea-0c69-b13a4340a01a
count(trials)

# ╔═╡ 1b778d18-bdaf-11ea-3ee4-0f3a16536b9f
function bernoulli_experiment(p, N=100)
	trials = [bernoulli(p) for i ∈ 1:N]
	return count(trials)
end

# ╔═╡ 45a5c6a4-bdaf-11ea-1812-359502135397
count(trials .== false) + count(trials)

# ╔═╡ 65bf3b8a-bdaf-11ea-1878-bb69724e1a7b
count(.!(trials)) # ! is "not" in Julia

# ╔═╡ 84062240-bdaf-11ea-3494-8d49391f642a
bernoulli_experiment(0.25)

# ╔═╡ a9b7c836-bdaf-11ea-1691-9d38af43f16a
bernoulli_experiment(0.25)

# ╔═╡ be407118-bdaf-11ea-2973-853fc1a266e1
md"A quantity that has different **outcomes** in different runs is called a **random variable**"

# ╔═╡ de5398ea-bdaf-11ea-1d81-5df4eb22e68b
md"## Monte Carlo simulation"

# ╔═╡ 73e79708-bdb0-11ea-123f-9fbb2c18dc12
md"Run the same random process a lot of times and look at the results"

# ╔═╡ 85f366ac-bdb0-11ea-16b7-2bec6b2f5356
begin
	local p = 0.25
	local N = 20 # number of trials
	local num_expts = 1000

	results = [bernoulli_experiment(p, N) for i in 1:num_expts]
end

# ╔═╡ af9e25e6-bdb0-11ea-23aa-7db6e35a4447
scatter(results, zcolor=results)

# ╔═╡ 87eaac94-bdb1-11ea-117a-ffa693e0b098
md" $X$ = number of heads

**Probability distribution** of a random variable $X$

$\text{Prob}(X = x)$
"

# ╔═╡ 976c1838-bdb1-11ea-30b1-f5c5eb8e2360
md"Probability that $X = x$ = proportion of time that the result was $x$"

# ╔═╡ f411c49a-bdb2-11ea-00b9-dfcc68181d37
md"## Counting"

# ╔═╡ 0922211a-bdb3-11ea-1e49-3518de406ace
md"Need to count how many of the trials come out as 5"

# ╔═╡ 286acb30-bdb3-11ea-1be2-c51955c53ab8
md"Need data structure where we store the counts:

- `Dict` (dictionary)
- `Vector`"

# ╔═╡ 50d0c656-bdb3-11ea-0a32-f7d63d08f64a
results

# ╔═╡ 99aab29c-bdb3-11ea-3dd0-57d1e619f762
maximum(results) #not 'max': max(1, 2, 3)

# ╔═╡ 2beb3cb0-bdb6-11ea-1451-05c1e3087165
minimum(results)

# ╔═╡ 6e8a4d8a-bdb7-11ea-0fd5-6d942ceb5283
l = maximum(results) + 1

# ╔═╡ 8e80e7fc-bdb7-11ea-1967-bbff2030d913
counts = zeros(Int, l)

# ╔═╡ 9b648744-bdb7-11ea-3e0e-7307a4084f9a
for score ∈ results
	counts[score + 1] += 1
end

# ╔═╡ e2635850-bdb7-11ea-21bc-5beadadfa675
counts

# ╔═╡ 7aecb5b0-bdba-11ea-026f-09ebc21a36a6
import Pkg; Pkg.add("LaTeXStrings")

# ╔═╡ 96b1bbb2-bdb8-11ea-3e81-1bca33878a99
begin
	plot(0:maximum(results), counts, m=:o)
	ylabel!("frequency of having n heads")
	xlabel!("n")
	#xlabel!("n")
end

# ╔═╡ 7090316e-bdba-11ea-27a6-dbc5ec194d02
function count_them(results)
	counts = zeros(Int, maximum(results) + 1)

	for score ∈ results
		counts[score + 1] += 1
	end

	return counts
end

# ╔═╡ c56f3320-bdbb-11ea-2904-6197a97efbbf
function run_experiments(p, N, num_expts=1000)
	results = [bernoulli_experiment(p, N) for i ∈ 1:num_expts]

	return results
end

# ╔═╡ 4c85f43c-bdbc-11ea-1ba1-810cadb75f6e
begin
	local data = count_them(run_experiments(0.25, 20, 10000))
	plot(0:length(data)-1, data, m=:o)
end

# ╔═╡ 8c2374a2-bdbc-11ea-12a9-eda80a7e4e21
md"## Probability distribution"

# ╔═╡ a00fa58a-bdbc-11ea-14a2-fd43cfedf7bb
md"Probability = relative frequency

Divide by the number of experiments"

# ╔═╡ bc8cf4ec-bdbc-11ea-33f3-d9aa8f26cd60
begin
	local data = run_experiments(0.25, 20, 10^5) |> count_them
	plot(0:length(data)-1, data ./ sum(data), m=:o)
end

# ╔═╡ 0839ca58-bdbd-11ea-2a3e-c99ad45c7121
begin
	num_expts = 10^5
	data = run_experiments(0.25, 20, num_expts) |> count_them
	probs = data .// num_expts
end

# ╔═╡ cb4d967a-bdbd-11ea-35c8-ad831d395195
sum(probs)

# ╔═╡ cf802b2c-bdbd-11ea-0d10-390f5ea9fe19
x = 3 // 4

# ╔═╡ fb862028-bdbd-11ea-3c26-e5e8ebd9630d
typeof(x)

# ╔═╡ 04180cb0-bdbe-11ea-2a64-a33a4f778d68
x + x

# ╔═╡ 08a5bf84-bdbe-11ea-18d4-dd7255a35b56
data

# ╔═╡ 741484a8-bdbe-11ea-3fd8-e55fe46723f8
begin
	local results = run_experiments(0.25, 20, 10^5)
	mean(results)
end

# ╔═╡ 5d7c00b8-bdbe-11ea-2b46-6d1117a42368
md"Expected mean value = N * p"

# ╔═╡ Cell order:
# ╟─fc08751c-bd79-11ea-072b-33a239fa0605
# ╟─f895f5da-bd7c-11ea-2250-efdcd9d49a38
# ╠═65f53264-bd7d-11ea-03cc-9513654c222c
# ╠═2dce78b0-bd7e-11ea-0381-d16883397ebd
# ╟─38ad903e-bd7e-11ea-08b6-f9a9a086260f
# ╠═c7c0a170-bd7f-11ea-3b68-a9ae0c8ccfde
# ╠═e40b10e0-bd7f-11ea-2883-bf485b88d6ad
# ╠═b9db886e-bd80-11ea-1347-5d5415a22dd5
# ╠═c427a4fc-bd80-11ea-365a-59fd43802b47
# ╠═92666d14-bd81-11ea-2ba7-e181ffa9f476
# ╠═0f3c4b1c-bd80-11ea-0fd0-df1c82c3ed1c
# ╠═71759792-bd81-11ea-1cbe-35ca8456b6a1
# ╟─7a610d46-bd81-11ea-0ff9-5b97563bce1e
# ╠═29e900e6-bd82-11ea-0344-2b2276b99422
# ╠═77489ff6-bd82-11ea-2d09-8737634cc579
# ╠═8cec500a-bd82-11ea-2fbb-03dec9457d7e
# ╠═a1c76a96-bd82-11ea-22a0-d9110840ecf3
# ╠═d39e0c64-bd82-11ea-188d-6f762d4dd688
# ╠═e9b93d52-bd82-11ea-15c9-699eeebaa243
# ╠═ed82a2b8-bd82-11ea-1d5c-8759b29ce1cc
# ╠═0682ca0c-bd83-11ea-0afa-8794eb7bd9b8
# ╟─be2d6b30-bd83-11ea-2683-49e64cbccb3c
# ╠═811488a0-bdab-11ea-0f9d-b1bc9f854e07
# ╟─89984d06-bdab-11ea-3ff4-55786243da18
# ╠═a8196c14-bdac-11ea-3b3d-f744310c289d
# ╟─cfac1f86-bdac-11ea-316c-a9abea7ed509
# ╠═46de61a4-bdad-11ea-3aaf-a5f9b96e533f
# ╠═8a20f6b6-bdad-11ea-0ef5-af788315bd82
# ╠═bc82b248-bdad-11ea-3ad5-019463cf2ef1
# ╠═c7552f0e-bdad-11ea-3fe8-fd3cf143ff30
# ╠═fe9569d2-bdad-11ea-0f72-09084f78f692
# ╠═38c6f580-bdae-11ea-2cbe-935e598ee04c
# ╠═3a406aea-bdae-11ea-0c69-b13a4340a01a
# ╠═1b778d18-bdaf-11ea-3ee4-0f3a16536b9f
# ╠═45a5c6a4-bdaf-11ea-1812-359502135397
# ╠═65bf3b8a-bdaf-11ea-1878-bb69724e1a7b
# ╠═84062240-bdaf-11ea-3494-8d49391f642a
# ╠═a9b7c836-bdaf-11ea-1691-9d38af43f16a
# ╟─be407118-bdaf-11ea-2973-853fc1a266e1
# ╟─de5398ea-bdaf-11ea-1d81-5df4eb22e68b
# ╟─73e79708-bdb0-11ea-123f-9fbb2c18dc12
# ╠═85f366ac-bdb0-11ea-16b7-2bec6b2f5356
# ╠═af9e25e6-bdb0-11ea-23aa-7db6e35a4447
# ╟─87eaac94-bdb1-11ea-117a-ffa693e0b098
# ╟─976c1838-bdb1-11ea-30b1-f5c5eb8e2360
# ╟─f411c49a-bdb2-11ea-00b9-dfcc68181d37
# ╟─0922211a-bdb3-11ea-1e49-3518de406ace
# ╟─286acb30-bdb3-11ea-1be2-c51955c53ab8
# ╠═50d0c656-bdb3-11ea-0a32-f7d63d08f64a
# ╠═99aab29c-bdb3-11ea-3dd0-57d1e619f762
# ╠═2beb3cb0-bdb6-11ea-1451-05c1e3087165
# ╠═6e8a4d8a-bdb7-11ea-0fd5-6d942ceb5283
# ╠═8e80e7fc-bdb7-11ea-1967-bbff2030d913
# ╠═9b648744-bdb7-11ea-3e0e-7307a4084f9a
# ╠═e2635850-bdb7-11ea-21bc-5beadadfa675
# ╠═7aecb5b0-bdba-11ea-026f-09ebc21a36a6
# ╠═7eab739e-bdba-11ea-3ff3-fd88b08946d3
# ╠═96b1bbb2-bdb8-11ea-3e81-1bca33878a99
# ╠═7090316e-bdba-11ea-27a6-dbc5ec194d02
# ╠═c56f3320-bdbb-11ea-2904-6197a97efbbf
# ╠═4c85f43c-bdbc-11ea-1ba1-810cadb75f6e
# ╟─8c2374a2-bdbc-11ea-12a9-eda80a7e4e21
# ╟─a00fa58a-bdbc-11ea-14a2-fd43cfedf7bb
# ╠═bc8cf4ec-bdbc-11ea-33f3-d9aa8f26cd60
# ╠═0839ca58-bdbd-11ea-2a3e-c99ad45c7121
# ╠═cb4d967a-bdbd-11ea-35c8-ad831d395195
# ╠═cf802b2c-bdbd-11ea-0d10-390f5ea9fe19
# ╠═fb862028-bdbd-11ea-3c26-e5e8ebd9630d
# ╠═04180cb0-bdbe-11ea-2a64-a33a4f778d68
# ╠═08a5bf84-bdbe-11ea-18d4-dd7255a35b56
# ╠═508fed7e-bdbe-11ea-1d75-8935a47d06d8
# ╠═741484a8-bdbe-11ea-3fd8-e55fe46723f8
# ╟─5d7c00b8-bdbe-11ea-2b46-6d1117a42368


import StatsBase: sample

# Model selection

# Taken from https://royalsocietypublishing.org/doi/pdf/10.1098/rspa.2017.0009
function AIC(k::Int64, X::AbstractArray, Y::AbstractArray; likelyhood = (X,Y) -> sum(abs2, X-Y))
    @assert size(X) == size(Y) "Dimensions of trajectories should be equal !"
    return 2*k - 2*log(likelyhood(X, Y))
end
# Taken from https://royalsocietypublishing.org/doi/pdf/10.1098/rspa.2017.0009
function AICC(k::Int64, X::AbstractArray, Y::AbstractArray; likelyhood = (X,Y) -> sum(abs2, X-Y))
    @assert size(X) == size(Y) "Dimensions of trajectories should be equal !"
    return AIC(k, X, Y, likelyhood = likelyhood)+ 2*(k+1)*(k+2)/(size(X)[2]-k-2)
end

# Double check on that
# Taken from https://www.immagic.com/eLibrary/ARCHIVES/GENERAL/WIKIPEDI/W120607B.pdf
function BIC(k::Int64, X::AbstractArray, Y::AbstractArray; likelyhood = (X,Y) -> sum(abs2, X-Y))
    @assert size(X) == size(Y) "Dimensions of trajectories should be equal !"
    return - 2*log(likelyhood(X, Y)) + k*log(size(X)[2])
end

# Optimal Shrinkage for data in presence of white noise
# See D. L. Donoho and M. Gavish, "The Optimal Hard Threshold for Singular
# Values is 4/sqrt(3)", http://arxiv.org/abs/1305.5870
# Code taken from https://github.com/erichson/optht

function optimal_svht(m::Int64, n::Int64; known_noise::Bool = false)
    @assert m/n > 0
    @assert m/n <= 1

    β = m/n
    ω = (8*β) / (β+1+sqrt(β^2+14β+1))
    c = sqrt(2*(β+1)+ω)

    if known_noise
        return c
    else
        median = median_marcenko_pastur(β)
        return c / sqrt(median)
    end
end

function marcenko_pastur_density(t, lower, upper, beta)
    sqrt((upper-t).*(t-lower))./(2π*beta*t)
end

function incremental_marcenko_pastur(x, beta, gamma)
    @assert beta <= 1
    upper = (1+sqrt(beta))^2
    lower = (1-sqrt(beta))^2

    @inline marcenko_pastur(x) = begin
        if (upper-x)*(x-lower) > 0
            return marcenko_pastur_density(x, lower, upper, beta)
        else
            return zero(eltype(x))
        end
    end

    if gamma ≈ zero(eltype(gamma))
        i, ϵ = quadgk(x->(x^gamma)*marcenko_pastur(x), x, upper)
        return i
    else
        i, ϵ = quadgk(x->marcenko_pastur(x), x, upper)
        return i
    end
end

function median_marcenko_pastur(beta)
    @assert 0 < beta <= 1
    upper = (1+sqrt(beta))^2
    lower = (1-sqrt(beta))^2
    change = true
    x = ones(eltype(upper), 5)
    y = similar(x)
    while change && (upper - lower > 1e-5)
        x = range(lower, upper, length = 5)
        for (i,xi) in enumerate(x)
            y[i] = one(eltype(x)) - incremental_marcenko_pastur(xi, beta, 0)
        end
        any(y .< 0.5) ? lower = maximum(x[y .< 0.5]) : change = false
        any(y .> 0.5) ? upper = minimum(x[y .> 0.5]) : change = false
    end
    return (lower+upper)/2
end

function optimal_shrinkage(X::AbstractArray{T, 2}) where T <: Number
    m,n = minimum(size(X)), maximum(size(X))
    U, S, V = svd(X)
    τ = optimal_svht(m,n)
    inds = S .>= τ*median(S)
    return U[:, inds]*Diagonal(S[inds])*V[:, inds]'
end

function optimal_shrinkage!(X::AbstractArray{T, 2}) where T <: Number
    m,n = minimum(size(X)), maximum(size(X))
    U, S, V = svd(X)
    τ = optimal_svht(m,n)
    inds = S .>= τ*median(S)
    X .= U[:, inds]*Diagonal(S[inds])*V[:, inds]'
    return
end

function savitzky_golay(x::AbstractVector{T}, windowSize::Integer, polyOrder::Integer; deriv::Integer=0, dt::Real=1.0, crop::Bool = true) where T <: Number
	# Polynomial smoothing with the Savitzky Golay filters
	# Adapted from: https://github.com/BBN-Q/Qlab.jl/blob/master/src/SavitskyGolay.jl
	# More information: https://pdfs.semanticscholar.org/066b/7534921b308925f6616480b4d2d2557943d1.pdf
	# Requires LinearAlgebra and DSP modules loaded.

	# Some error checking
	@assert isodd(windowSize) "Window size must be an odd integer."
	@assert polyOrder < windowSize "Polynomial order must be less than window size."

	# Calculate filter coefficients
	filterCoeffs = calculate_filterCoeffs(windowSize, polyOrder, deriv, dt)

	# Pad the signal with the endpoints and convolve with filter
	halfWindow = Int(ceil((windowSize - 1)/2))
	paddedX = [x[1]*ones(halfWindow); x; x[end]*ones(halfWindow)]
	y = conv(filterCoeffs[end:-1:1], paddedX)

	if !crop
		# Return the valid midsection
		return y[2*halfWindow+1:end-2*halfWindow]
	else
		# Return cropped data. Excluding borders, where the estimation is less accurate
		return x[halfWindow+2:end-halfWindow-1], y[3*halfWindow+2:end-3*halfWindow-1]
	end
end

function savitzky_golay(x::AbstractMatrix{T}, windowSize::Integer, polyOrder::Integer; deriv::Integer=0, dt::Real=1.0, crop::Bool = true) where T <: Number
	# Polynomial smoothing with the Savitzky Golay filters
	# Adapted from: https://github.com/BBN-Q/Qlab.jl/blob/master/src/SavitskyGolay.jl
	# More information: https://pdfs.semanticscholar.org/066b/7534921b308925f6616480b4d2d2557943d1.pdf
	# Requires LinearAlgebra and DSP modules loaded.

	# Some error checking
	@assert isodd(windowSize) "Window size must be an odd integer."
	@assert polyOrder < windowSize "Polynomial order must be less than window size."

	# Calculate filter coefficients
	filterCoeffs = calculate_filterCoeffs(windowSize, polyOrder, deriv, dt)

	# Apply filter to each component
	halfWindow = Int(ceil((windowSize - 1)/2))

	if !crop
		y = similar(x)
		for (i, xi) in enumerate(eachrow(x))
			paddedX = [xi[1]*ones(halfWindow); xi; xi[end]*ones(halfWindow)]
			y₀ = conv(filterCoeffs[end:-1:1], paddedX)
			y[i,:] = y₀[2*halfWindow+1:end-2*halfWindow]
		end
		return y
	else
		cropped_x = x[:,halfWindow+2:end-halfWindow-1]
		y = similar(cropped_x)
		for (i, xi) in enumerate(eachrow(x))
			paddedX = [xi[1]*ones(halfWindow); xi; xi[end]*ones(halfWindow)]
			y₀ = conv(filterCoeffs[end:-1:1], paddedX)
			y[i,:] = y₀[3*halfWindow+2:end-3*halfWindow-1]
		end
		return cropped_x, y
	end
end

function calculate_filterCoeffs(windowSize::Integer, polyOrder::Integer, deriv::Integer, dt::Real)
	# Some error checking
	@assert isodd(windowSize) "Window size must be an odd integer."
	@assert polyOrder < windowSize "Polynomial order must be less than window size."

	# Form the design matrix A
	halfWindow = Int(ceil((windowSize - 1)/2))
	A = zeros(windowSize, polyOrder+1)
	for order = 0:polyOrder
		A[:, order+1] = (-halfWindow:halfWindow).^(order)
	end

	# Compute the required column of the inverse of A'*A
	# and calculate filter coefficients
	ei = zeros(polyOrder+1)
	ei[deriv+1] = 1.0
	inv_col = (A'*A) \ ei
	return A*inv_col * factorial(deriv) ./(dt^deriv)
end

@inline function burst_sampling(x::AbstractArray, samplesize::Int64, bursts::Int64)
    @assert size(x)[end] >= samplesize*bursts "Length of data array too small for subsampling of size $size!"
    inds = sample(1:size(x)[end]-samplesize, bursts, replace = false)
    inds = sort(unique(vcat([collect(i:i+samplesize) for i in inds]...)))
    return resample(x, inds)
end

@inline function burst_sampling(x::AbstractArray, y::AbstractArray, samplesize::Int64, bursts::Int64)
    @assert size(x)[end] >= samplesize*bursts "Length of data array too small for subsampling of size $size!"
    @assert size(x)[end] == size(y)[end]
    inds = sample(1:size(x)[end]-samplesize, bursts, replace = false)
    inds = sort(unique(vcat([collect(i:i+samplesize) for i in inds]...)))
    return resample(x, inds), resample(y, inds)
end

@inline function burst_sampling(x::AbstractArray, t::AbstractVector, period::T, bursts::Int64) where T <: AbstractFloat
    @assert period > zero(typeof(period)) "Sampling period has to be positive."
    @assert size(x)[end] == size(t)[end] "Provide consistent data."
    @assert bursts >= 1 "Number of bursts has to be positive."
    @assert t[end]-t[1]>= period*bursts "Bursting impossible. Please provide more data or reduce bursts."
    t_ids = zero(eltype(t)) .<= t .- period  .<= t[end] .- 2*period
    samplesize = Int64(floor(period/(t[end]-t[1])*length(t)))
    inds = sample(collect(1:length(t))[t_ids], bursts, replace = false)
    inds = sort(unique(vcat([collect(i:i+samplesize) for i in inds]...)))
    return resample(x, inds), resample(t, inds)
end

@inline function subsample(x::AbstractVector, frequency::Int64)
    @assert frequency > 1
    return x[1:frequency:end]
end

@inline function subsample(x::AbstractArray, frequency::Int64)
    @assert frequency > 1
    return x[:, 1:frequency:end]
end

@inline function subsample(x::AbstractArray, t::AbstractVector, period::T) where T <: AbstractFloat
    @assert period > zero(typeof(period)) "Sampling period has to be positive."
    @assert size(x)[end] == size(t)[end] "Provide consistent data."
    @assert t[end]-t[1]>= period "Subsampling impossible. Sampling period exceeds time window."
    idx = Int64[1]
    t_now = t[1]
    @inbounds for (i, t_current) in enumerate(t)
        if t_current - t_now >= period
            push!(idx, i)
            t_now = t_current
        end
    end
    return resample(x, idx), resample(t, idx)
end

@inline function resample(x::AbstractArray{T,1}, indx::AbstractArray{Int64}) where T <: Number
    @assert maximum(indx) <= length(x)
    @assert minimum(indx) >= 1
    return x[indx]
end

@inline function resample(x::AbstractArray{T,2}, indx::AbstractArray{Int64}) where T <: Number
    @assert maximum(indx) <= size(x, 2)
    @assert minimum(indx) >= 1
    return x[:, indx]
end


using DifferentialEquations
using SimpleDiffEq
using DataFrames
using StatsPlots
using BenchmarkTools

@inline function rate_to_proportion(r::Float64,t::Float64)
    1-exp(-r*t)
end;

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

δt = 0.1
nsteps = 400
tmax = nsteps*δt
tspan = (0.0,nsteps)
t = 0.0:δt:tmax;

u0 = [990.0,10.0,0.0];

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

prob_map = DiscreteProblem(sir_map!,u0,tspan,p);

sol_map = solve(prob_map,solver=FunctionMap);

df_map = DataFrame(sol_map')
df_map[!,:t] = t;

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

@benchmark solve(prob_map,solver=FunctionMap)


struct GeneralizedPoisson <: Distributions.DiscreteUnivariateDistribution
    θ::Float64
    λ::Float64
end

function Distributions.pdf(G::GeneralizedPoisson, x::Int64)
    G.θ*(G.θ+G.λ*x)^(x-1)*exp(-G.θ - G.λ*x)/factorial(x)
end

function Distributions.logpdf(G::GeneralizedPoisson, x::Int64)
    a = log(G.θ)
    a += (x-1)*log(G.θ+G.λ*x)
    a -= G.θ
    a -= G.λ*x
    a -= lgamma(x+1)
    a
end

Distributions.mean(G::GeneralizedPoisson) = G.θ/(1-G.λ)
Distributions.var(G::GeneralizedPoisson) = G.θ/(1-G.λ)^3


using Pkg
Pkg.activate(pwd())

# # Arrays
# ## Vectors

v = [1, 2, 3, 4, 5, 6, 7, 8] # or equivalently v = [1; 2; 3; 4; ...]

#+

typeof(v)
ndims(v)
eltype(v)
size(v)
length(v)

#+

v[3]
v[begin] # the first element
v[end] # the last element
v[[2, 3]]

#+

range(1; stop = 10, step = 2) # or equivalently range(1, 10; step = 2)
1:2:10

#+

v[1:3] # the first three elements
v[1:2:end] # select all elements with odd index
v[:] # all elements

#+

v = [1,2,3]

append!(v, 4)
append!(v, [5,6])
append!(v, 7:8)
append!(v, 3.0)
append!(v, 3.1415)

#+

isinteger(3.0)

#+

v = Float64[1, 2, 3]
append!(v, 3.1415)

#+

v = [1, 2, 3, 4]
v[2] = 4
v

#+

v[3:4] .= 11
v

# ### Exercise:
# Create a vector of positive integers that contains all odd numbers smaller than `10`.
# Then change the first element to `4` and the last two elements to `1`.
#
# ---
# ### Solution:

v = [1,3,5,7,9]

collect(1:2:9)
Vector(1:2:9)

v[1] = 4
v[end-1:end] .= 1
v

# ---
#
# ## Matrices

m = [1  2  3  4; 5  6  7  8]

typeof(m)
eltype(m)
ndims(m)
size(m)
length(m)

#+

m[1] # the first element, equivalent to m[begin]
m[2] # the second element
m[end-1] # the last element
m[1, 2]
m[1, [2, 3]] # the second and third element in the first row
m[1:3] # the first three elements according to linear indexing
m[:, 1:3] # the first three columns
m[1, :] # the first row
m[:] # all elements

#+

hcat(m, m)
vcat(m, m)

#+

cat(m, m; dims = 2) # equivalent to hcat(m, m)
cat(m, m; dims = 1) # equivalent to vcat(m, m)

#+

v = [11, 12]
vcat(m, v)

# ### Exercise:
# Create two vectors: vector of all odd positive integers smaller than `10` and vector of
# all even positive integers smaller than `10`. Then concatenate these two vectors
# horizontally and fill the third row with `4`.
#
# ---
# ### Solution:

v1 = collect(1:2:9)
v2 = collect(2:2:10)

M = hcat(v1, v2)
M[3,:] .= 4
M

# ---
#
# ## `N`-dimensional arrays

A = zeros(3, 5, 2) # equivalent to A = zeros((3, 5, 2))
B = zeros(Int64, 3, 5, 2)  # equivalent to B = zeros(Int64, (3, 5, 2))

typeof(A)
eltype(A)
ndims(A)
size(A)
length(A)

#+

B[1] = 1 # assign 1 to the first element
B[1, 2, 2] = 2 # assign 2 to the element at position (1,2,2)
B[2,:,1] .= 4

B

#+

ones(Float32, 2, 3, 1)
fill(1.234, 2, 2)

# ### Exercise:
# Create three matrices with the following properties:
# - Matrix `A` is of size `2x3`, and all its elements equal 0.
# - Matrix `B` is of size `2x3x1`, and all its elements equal 1.
# - Matrix `C` is of size `2x3`, and all its elements equal 2.
# Concatenate these three matrices along the third dimension.
#
# **Hint:** use the `cat` function and the keyword `dims`.
#
# ---
# ### Solution:

A = zeros(2, 3)
B = ones(2, 3, 1)
C = fill(2, 2, 3)

cat(A, B, C; dims = 3)

# ---
#
# ## Broadcasting

a = [1,2,3] # column vector

a .-= 4 # from each element of vector subtracts 4
a -= 1

#+

abs.(a)
sum(exp.(sqrt.(abs.(a .- 1)))./2)

#+

a = [1,2,3] # column vector
b = [4,5,6] # column vector

a * b
a' * b
a * b'
a .* b

# ### Exercise:
# Construct a matrix whose elements are given by the following formula
# $$
# A_{i, j} = \frac{1}{2}\exp\{(B_{i, j} + 1)^2\}, \quad i \in \{1, 2\}, \; j \in
# \{1, 2, 3\}
# $$
# where the matrix $B$ is defined by

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

#
# ---
# ### Solution:

A = exp.((B .+ 1) .^ 2) ./ 2

A = @. exp((B + 1) ^ 2) / 2

A = zeros(2, 3);
for i in 1:length(A)
    A[i] = exp((B[i] + 1)^2)/2
end
A

# ---
#
# ## Views

A = [1 2 3; 4 5 6]
B = A

B[2] = 42

A == B

#+

C = copy(A)
C[4] = 10

A == C

#+

D = A[1:2, 1:2]
D[1] = 15

D == A[1:2, 1:2]

#+

E = view(A, 1:2, 1:2)
E = @view A[1:2, 1:2]

E[4] = 78

E == A[1:2, 1:2]

#+

typeof(E)

#+

A = [1 2 3; 4 5 6]
A_view = @view A[:, :]

sum(A)
sum(A_view)
minimum(A; dims = 1)
minimum(A_view; dims = 1)

#+

A = [1 2 3; 4 5 6];

sum(exp.(sqrt.(abs.(@view(A[1, :]) .- @view(A[2, :]))))./2)
@views sum(exp.(sqrt.(abs.(A[1, :] .- A[2, :])))./2)


# # Binary (or 0-1) knapsack problem
# Given a knapsack of some capacity $C$ and $n$ objects with object $i$ having weight $w_i$ and profit $p_i$, the goal is to choose some subset of the objects that can fit in the knapsack (i.e. the sum of their weights is no more than $C$) while maximizing profit.
#
# This can be formulated as a mixed-integer program as:
#
# $$
# \begin{array}{ll}
#   \text{maximize} & x' p \\
#     \text{subject to} & x \in \{0, 1\} \\
#   & w' x \leq C \\
# \end{array}
# $$
#
# where $x$ is a vector is size $n$ where $x_i$ is one if we chose to keep the object in the knapsack, 0 otherwise.

## Data taken from http://people.sc.fsu.edu/~jburkardt/datasets/knapsack_01/knapsack_01.html
w = [23; 31; 29; 44; 53; 38; 63; 85; 89; 82]
C = 165
p =  [92; 57; 49; 68; 60; 43; 67; 84; 87; 72];
n = length(w)

#-

using Convex, GLPK
x = Variable(n, :Bin)
problem = maximize(dot(p, x), dot(w, x) <= C)
solve!(problem, GLPK.Optimizer)
evaluate(x)


function hessian{T<:AbstractFloat}(f::Function,x::Array{T,1})

  m = length(x)

	e = eps(T)^(1/4)*maxabs([x;one(T)])
	dh = eye(m)*e
  hess = Array(T,m,m)

  for i = 1:m
		hess[i,i] = (-f(x+2*dh[:,i])+16*f(x+dh[:,i])-30*f(x)+16*f(x-dh[:,i])-f(x-2*dh[:,i]))/(12*e^2)
  end

    for i = 1:m-1
      for j = i+1:m
			  hess[i,j] = (f(x+dh[:,i]+dh[:,j])+f(x-dh[:,i]-dh[:,j])-f(x-dh[:,i]+dh[:,j])-f(x+dh[:,i]-dh[:,j]))/(4*e^2)
        hess[j,i] = hess[i,j]
      end
    end

    return hess

end


"""
    MaternKernel(; ν = 1.0)

The matern kernel is a Mercer kernel given by the formula:
```
    κ(x,y) = 2^{1-ν}/Γ(ν)*(√(2ν)‖x-y‖)^ν K_ν(√(2ν)‖x-y‖)
```
For `ν=n+1/2, n=0,1,2,...` it can be simplified and you should instead use
[`ExponentialKernel`](@ref) for `n=0`, [`Matern32Kernel`](@ref), for `n=1`,
[`Matern52Kernel`](@ref) for `n=2` and [`SqExponentialKernel`](@ref) for `n=∞`.
"""
struct MaternKernel{Tν<:Real} <: SimpleKernel
    ν::Vector{Tν}
    function MaternKernel(;nu::T=1.5, ν::T=nu) where {T<:Real}
        @check_args(MaternKernel, ν, ν > zero(T), "ν > 0")
        return new{T}([ν])
    end
end

@functor MaternKernel

@inline function kappa(κ::MaternKernel, d::Real)
    result = _matern(first(κ.ν), d)
    return ifelse(iszero(d), one(result), result)
end

function _matern(ν::Real, d::Real)
    y = sqrt(2ν) * d
    return exp((one(d) - ν) * logtwo - loggamma(ν) + ν * log(y) + log(besselk(ν, y)))
end

metric(::MaternKernel) = Euclidean()

Base.show(io::IO, κ::MaternKernel) = print(io, "Matern Kernel (ν = ", first(κ.ν), ")")

"""
    Matern32Kernel()

The matern 3/2 kernel is a Mercer kernel given by the formula:
```
    κ(x,y) = (1+√(3)‖x-y‖)exp(-√(3)‖x-y‖)
```
"""
struct Matern32Kernel <: SimpleKernel end

kappa(κ::Matern32Kernel, d::Real) = (1 + sqrt(3) * d) * exp(-sqrt(3) * d)

metric(::Matern32Kernel) = Euclidean()

Base.show(io::IO, ::Matern32Kernel) = print(io, "Matern 3/2 Kernel")

"""
    Matern52Kernel()

The matern 5/2 kernel is a Mercer kernel given by the formula:
```
    κ(x,y) = (1+√(5)‖x-y‖ + 5/3‖x-y‖^2)exp(-√(5)‖x-y‖)
```
"""
struct Matern52Kernel <: SimpleKernel end

kappa(κ::Matern52Kernel, d::Real) = (1 + sqrt(5) * d + 5 * d^2 / 3) * exp(-sqrt(5) * d)

metric(::Matern52Kernel) = Euclidean()

Base.show(io::IO, ::Matern52Kernel) = print(io, "Matern 5/2 Kernel")


#=
Generate Symbolic Jacobian with ModelingToolkit. The testing equations are from Matlab tutorial:
- Jacobian and gradient: https://www.mathworks.com/help/symbolic/jacobian.html
- Hessian: https://www.mathworks.com/help/symbolic/hessian.html
=#

using ModelingToolkit

@variables x, y, z

# Jacobian: Vector{1, N} → Matrix{N,M} where M is Number of variables
jac = ModelingToolkit.jacobian([x*y*z, y^2, x + z], [x, y, z])

# Gradient: Equation → Vector
grad = ModelingToolkit.gradient(2*x + 3*y + 4*z, [x, y, z])

# Hessian
hess = ModelingToolkit.hessian(x*y+2*z*x, [x,y,z])

# Hessian: Jacobian of Gradient
hess_ = ModelingToolkit.jacobian(ModelingToolkit.gradient(x*y+2*z*x, [x,y,z]), [x,y,z])


# logarithmic barrier function

immutable LogBarrier{T <: Real} <: ProximableFunction
  a::T
  b::T
  mu::T
  function LogBarrier(a::T, b::T, mu::T)
    if mu <= 0
      error("parameter mu must be positive")
    else
      new(a, b, mu)
    end
  end
end

"""
  LogBarrier(a::Real=1.0, b::Real=0.0, mu::Real=1.0)

Returns the function `g(x) = -mu*sum(log(a*x_i+b), i=1,...,n)`.
"""

LogBarrier{T <: Real}(a::T=1.0, b::T=0.0, mu::T=1.0) = LogBarrier{T}(a, b, mu)

function (f::LogBarrier){T <: Real}(x::AbstractArray{T,1})
  sumf = 0.0
  v = 0.0
  for i in eachindex(x)
    v = f.a*x[i]+f.b
    if v <= 0.0
      return +Inf
    end
    sumf += log(v)
  end
  return -f.mu*sumf
end

function prox!{T <: Real}(f::LogBarrier, x::AbstractArray{T,1}, y::AbstractArray{T}, gamma::Real=1.0)
  par = 4*gamma*f.mu*f.a*f.a
  sumf = 0.0
  z = 0.0
  v = 0.0
  for i in eachindex(x)
    z = f.a*x[i] + f.b
    v = (z + sqrt(z*z + par))/2
    y[i] = (v - f.b)/f.a
    sumf += log(v)
  end
  return -f.mu*sumf
end

fun_name(f::LogBarrier) = "logarithmic barrier"
fun_dom(f::LogBarrier) = "AbstractArray{Real}"
fun_expr(f::LogBarrier) = "x ↦ -μ * sum( log(a*x_i+b), i=1,...,n )"
fun_params(f::LogBarrier) = "a = $(f.a), b = $(f.b), μ = $(f.mu)"

function prox_naive{T <: Real}(f::LogBarrier, x::AbstractArray{T,1}, gamma::Real=1.0)
  asqr = f.a*f.a
  z = f.a*x + f.b
  y = ((z + sqrt.(z.*z + 4*gamma*f.mu*asqr))/2 - f.b)/f.a
  fy = -f.mu * sum(log.(f.a*y+f.b))
  return y, fy
end


module problem124

description = """
Ordered radicals
The radical of n, rad(n), is the product of the distinct prime factors of n. For example, 504 = 23 × 32 × 7, so rad(504) = 2 × 3 × 7 = 42.
Let E(k) be the kth element in the sorted n column; for example, E(4) = 8 and E(6) = 9.

If rad(n) is sorted for 1 ≤ n ≤ 100000, find E(10000).
""

# Calculate rad for all values between 1 and n using algorithm similar to sieve of eratosthenes
function radsieve(n)
  rad = ones(Int, n)
  for i = 2:n
    if rad[i] == 1
      for j = i:i:n
        rad[j] *= i
      end
    end
  end
  return rad
end

function solve(n=100000, k=10000)
  radsorted = sort(collect(enumerate(radsieve(n))), by = pr -> pr[2], alg=MergeSort)
  return radsorted[k][1]
end

using Base.Test

@test radsieve(10) == [1, 2, 3, 2, 5, 6, 7, 2, 3, 10]
@test solve(10, 4) == 8
@test solve(10, 6) == 9

end


using ModelingToolkit, MethodOfLines, LinearAlgebra, OrdinaryDiffEq
using ModelingToolkit: operation, istree, arguments
using DomainSets
using NonlinearSolve

# # Define some variables
@testset "Heat Equation 1D 2 variables" begin
    @parameters t x
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2
    eqs = [Dt(u(t, x)) ~ Dxx(u(t, x)),
        Dt(v(t, x)) ~ Dxx(v(t, x))]
    bcs = [u(0, x) ~ -x * (x - 1) * sin(x),
        v(0, x) ~ -x * (x - 1) * sin(x),
        u(t, 0) ~ 0.0, u(t, 1) ~ 0.0,
        v(t, 0) ~ 0.0, v(t, 1) ~ 0.0]

    domains = [t ∈ Interval(0.0, 1.0),
        x ∈ Interval(0.0, 1.0)]

    @named pdesys = PDESystem(eqs, bcs, domains, [t, x], [u(t, x), v(t, x)])
    discretization = MOLFiniteDifference([x => 0.1], t; grid_align=edge_align)
    prob = discretize(pdesys, discretization) # This gives an ODEProblem since it's time-dependent
    sol = solve(prob, Tsit5())
end

@testset "Heat Equation 2D 1 variable" begin
    @parameters t x y
    @variables u(..)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2
    Dt = Differential(t)
    t_min = 0.0
    t_max = 2.0
    x_min = 0.0
    x_max = 2.0
    y_min = 0.0
    y_max = 2.0

    # 3D PDE
    eq = Dt(u(t, x, y)) ~ Dxx(u(t, x, y)) + Dyy(u(t, x, y))

    analytic_sol_func(t, x, y) = exp(x + y) * cos(x + y + 4t)
    # Initial and boundary conditions
    bcs = [u(t_min, x, y) ~ analytic_sol_func(t_min, x, y),
        u(t, x_min, y) ~ analytic_sol_func(t, x_min, y),
        u(t, x_max, y) ~ analytic_sol_func(t, x_max, y),
        u(t, x, y_min) ~ analytic_sol_func(t, x, y_min),
        u(t, x, y_max) ~ analytic_sol_func(t, x, y_max)]

    # Space and time domains
    domains = [t ∈ Interval(t_min, t_max),
        x ∈ Interval(x_min, x_max),
        y ∈ Interval(y_min, y_max)]
    @named pdesys = PDESystem([eq], bcs, domains, [t, x, y], [u(t, x, y)])
    # Method of lines discretization
    dx = 0.1
    dy = 0.2
    discretization = MOLFiniteDifference([x => dx, y => dy], t)
    prob = ModelingToolkit.discretize(pdesys, discretization)
    sol = solve(prob, Tsit5())
end

# Diffusion in a sphere
@testset "Spherical Diffusion" begin

    @parameters t r
    @variables u(..)
    Dt = Differential(t)
    Dr = Differential(r)
    Drr = Dr^2
    eq = Dt(u(t, r)) ~ 1 / r^2 * Dr(r^2 * Dr(u(t, r)))
    bcs = [u(0, r) ~ -r * (r - 1) * sin(r),
        Dr(u(t, 0)) ~ 0.0, u(t, 1) ~ sin(1)]

    domains = [t ∈ Interval(0.0, 1.0),
        r ∈ Interval(0.0, 1.0)]

    @named pdesys = PDESystem(eq, bcs, domains, [t, r], [u(t, r)])
    discretization = MOLFiniteDifference([r => 0.1], t)
    prob = discretize(pdesys, discretization) # This gives an ODEProblem since it's time-dependent
    sol = solve(prob, Tsit5())
end

@testset "RHS = 0" begin
    @parameters x, t
    @variables u(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dx2 = Differential(x)^2
    Dx3 = Differential(x)^3
    Dx4 = Differential(x)^4

    α = 1
    β = 4
    γ = 1
    eq = Dt(u(x, t)) + u(x, t) * Dx(u(x, t)) + α * Dx2(u(x, t)) + β * Dx3(u(x, t)) + γ * Dx4(u(x, t)) ~ 0

    du(x, t; z=-x / 2 + t) = 15 / 2 * (tanh(z) + 1) * (3 * tanh(z) - 1) * sech(z)^2

    bcs = [u(x, 0) ~ x^2,
        Dx(u(-10, t)) ~ du(-10, t),
        Dx(u(10, t)) ~ du(10, t)]

    # Space and time domains
    domains = [x ∈ Interval(-10.0, 10.0),
        t ∈ Interval(0.0, 0.5)]
    # Discretization
    dx = 0.4
    dt = 0.2

    discretization = MOLFiniteDifference([x => dx], t; approx_order=4)
    @named pdesys = PDESystem(eq, bcs, domains, [x, t], [u(x, t)])
    prob = discretize(pdesys, discretization)
end

@testset "Wave Equation" begin
    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    Dt = Differential(t)             # Required in ICs and equation
    Dtt = Differential(t)^2           # required in equation
    Dxxxx = Differential(x)^4       # required in equation
    Dxx = Differential(x)^2         # required in BCs
    # some parameters
    EI = 291.6667
    m = 1.3850
    c = 1.0
    p = 0.01
    L = 2.0

    # 1D PDE and boundaru conditions
    eq = m * Dtt(u(t, x)) + c * Dt(u(t, x)) + EI * Dxxxx(u(t, x)) ~ 0

    ic_bc = [u(0, x) ~ (p * x * (x^3 + L^3 - 2 * L * x^2) / (24 * EI)), #for all 0 < u < L
        Dt(u(0, x)) ~ 0.0,        # for all 0 < u < L
        u(t, 0) ~ 0.0,            # for all t > 0,, displacement zero at u=0
        u(t, 2) ~ 0.0,            # for all t > 0,, displacement zero at u=L
        Dxx(u(t, 0)) ~ 0.0,       # for all t > 0,, curvature zero at u=0
        Dxx(u(t, 2)) ~ 0.0]       # for all t > 0,, curvature zero at u=L

    # Space and time domains
    domains = [t ∈ Interval(0.0, 2.0),
        x ∈ Interval(0.0, L)]

    dt = 0.1   # dt related to saving the data.. not actual dt
    # PDE sustem
    @named pdesys = PDESystem(eq, ic_bc, domains, [t, x], [u(t, x)])

    # Method of lines discretization
    dx = 0.1
    order = 2
    discretization = MOLFiniteDifference([x => dx, t => dt], approx_order=order)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys, discretization)

    # Solve the ODE problem
    sol = NonlinearSolve.solve(prob, NewtonRaphson())
end

@testset "Rearranged Robin" begin
    @parameters x t

    @variables c(..)

    ∂t = Differential(t)
    ∂x = Differential(x)
    ∂²x = Differential(x)^2

    D₀ = 1.0
    R = 0.5
    cₑ = 2.5
    ℓ = 2.0
    α = 1 / 3

    diff_eq = ∂t(c(x, t)) ~ ∂x((D₀ + α * c(x, t)) * ∂x(c(x, t)))

    bcs = [c(x, 0) ~ 0.0,         # initial condition
        R * (D₀ + α * c(0, t)) * ∂x(c(0, t)) ~ c(0, t) - cₑ, # Robin
        ∂x(c(ℓ, t)) ~ 0.0]   # no flux

    domains = [t ∈ Interval(0.0, 10.0),
        x ∈ Interval(0.0, ℓ)]

    @named pdesys = PDESystem(diff_eq, bcs, domains, [t, x], [c(x, t)])

    # Method of lines discretization
    dx = 0.1
    order = 2
    discretization = MOLFiniteDifference([x => dx], t, approx_order=order)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys, discretization)

    sol = solve(prob, Rodas4())
end

@testset "Array u" begin
    # Dependencies
    N = 6 # number of dependent variables

    # Variables, parameters, and derivatives
    @parameters x
    @variables u[1:N](..)
    Dx = Differential(x)
    Dxx = Differential(x)^2

    # Domain edges
    x_min = 0.0
    x_max = 1.0

    # Discretization parameters
    dx = 0.1
    order = 2

    #u = collect(u)

    # Equations
    eqs = Vector{ModelingToolkit.Equation}(undef, N)
    for i = 1:N
        eqs[i] = Dxx(u[i](x)) ~ u[i](x)
    end

    # Initial and boundary conditions
    bcs = Vector{ModelingToolkit.Equation}(undef, 2 * N)
    for i = 1:N
        bcs[i] = Dx(u[i](x_min)) ~ 0.0
    end

    for i = 1:N
        bcs[i+N] = u[i](x_max) ~ rand()
    end

    # Space and time domains
    domains = [x ∈ Interval(x_min, x_max)]

    # PDE system
    @named pdesys = PDESystem(eqs, bcs, domains, [x], collect([u[i](x) for i = 1:N]))

    # Method of lines discretization
    discretization = MOLFiniteDifference([x => dx], nothing, approx_order=order)
    prob = ModelingToolkit.discretize(pdesys, discretization)

    # # Solution of the ODE system
    sol = NonlinearSolve.solve(prob, NewtonRaphson())
end

@testset "Non-Uniform Grid" begin
    @parameters x y t
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dy = Differential(y)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2

    ∇²(u) = Dxx(u) + Dyy(u)

    brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0

    x_min = y_min = t_min = 0.0
    x_max = y_max = 1.0
    t_max = 11.5

    α = 10.0

    u0(x, y, t) = 22(y * (1 - y))^(3 / 2)
    v0(x, y, t) = 27(x * (1 - x))^(3 / 2)

    eq = [Dt(u(x, y, t)) ~ 1.0 + Dy(v(x, y, t)) * u(x, y, t)^2 - 4.4 * u(x, y, t) + α * ∇²(u(x, y, t)) + brusselator_f(x, y, t),
        Dt(v(x, y, t)) ~ -3.4 * Dx(u(x, y, t)) - v(x, y, t) * u(x, y, t)^2 + α * ∇²(v(x, y, t))]

    domains = [x ∈ Interval(x_min, x_max),
        y ∈ Interval(y_min, y_max),
        t ∈ Interval(t_min, t_max)]

    bcs = [u(x, y, 0) ~ u0(x, y, 0),
        u(0, y, t) ~ 0.0,
        u(1, y, t) ~ 0.0,
        u(x, 0, t) ~ 0.0,
        u(x, 1, t) ~ 0.0,
        v(x, y, 0) ~ v0(x, y, 0),
        v(0, y, t) ~ 0.0,
        v(1, y, t) ~ 0.0,
        v(x, 0, t) ~ 0.0,
        v(x, 1, t) ~ 0.0]

    @named pdesys = PDESystem(eq, bcs, domains, [x, y, t], [u(x, y, t), v(x, y, t)])

    # Method of lines discretization
    N = 5

    dx = [0.0, cumsum(rand(0.05:0.01:0.2, 5))...]
    if dx[end] < 1.0
        push!(dx, 1.0)
    end

    dy = [0.0, cumsum(rand(0.05:0.01:0.2, 5))...]
    if dy[end] < 1.0
        push!(dy, 1.0)
    end

    order = 2

    discretization = MOLFiniteDifference([x => dx, y => dy], t, approx_order=order)

    prob = discretize(pdesys, discretization)

    sol = solve(prob, TRBDF2())

end

@testset "2D variable connected to 1D variable at boundary #33" begin
    @parameters t x r
    @variables u(..) v(..)
    Dt = Differential(t)
    Dx = Differential(x)
    Dxx = Differential(x)^2
    Dr = Differential(r)
    Drr = Differential(r)^2

    s = u(t, x) + v(t, x, 1)

    eqs = [Dt(u(t, x)) ~ Dxx(u(t, x)) + s,
        Dt(v(t, x, r)) ~ Drr(v(t, x, r))]
    bcs = [u(0, x) ~ 1,
        v(0, x, r) ~ 1,
        Dx(u(t, 0)) ~ 0.0, Dx(u(t, 1)) ~ 0.0,
        Dr(v(t, x, 0)) ~ 0.0, Dr(v(t, x, 1)) ~ s]

    domains = [t ∈ Interval(0.0, 1.0),
        x ∈ Interval(0.0, 1.0),
        r ∈ Interval(0.0, 1.0)]

    @named pdesys = PDESystem(eqs, bcs, domains, [t, x, r], [u(t, x), v(t, x, r)])

    # Method of lines discretization
    dx = 0.1
    dr = 0.1
    order = 2
    discretization = MOLFiniteDifference([x => dx, r => dr], t, approx_order=order)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys, discretization)

    sol = solve(prob, Tsit5())
end

@testset "Testing discretization of varied systems" begin
    @parameters x t

    @variables c(..)

    ∂t = Differential(t)
    ∂x = Differential(x)
    ∂²x = Differential(x)^2

    D₀ = 1.5
    α = 0.15
    χ = 1.2
    R = 0.1
    cₑ = 2.0
    ℓ = 1.0
    Δx = 0.1

    bcs = [
        # initial condition
        c(x, 0) ~ 0.0,
        # Robin BC
        ∂x(c(0.0, t)) / (1 + exp(α * (c(0.0, t) - χ))) * R * D₀ + cₑ - c(0.0, t) ~ 0.0,
        # no flux BC
        ∂x(c(ℓ, t)) ~ 0.0]

    # define space-time plane
    domains = [x ∈ Interval(0.0, ℓ), t ∈ Interval(0.0, 5.0)]

    @testset "Test 01: ∂t(c(x, t)) ~ ∂x(D * ∂x(c(x, t)))" begin
        D = D₀ / (1.0 + exp(α * (c(x, t) - χ)))
        diff_eq = ∂t(c(x, t)) ~ ∂x(D * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 02: ∂t(c(x, t)) ~ ∂x(D * ∂x(c(x, t)))" begin
        D = 1.0 / (1.0 + exp(α * (c(x, t) - χ)))
        diff_eq = ∂t(c(x, t)) ~ ∂x(D * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 03: ∂t(c(x, t)) ~ ∂x(1.0 / (1.0/D₀ + exp(α * (c(x, t) - χ))/D₀) * ∂x(c(x, t)))" begin
        diff_eq = ∂t(c(x, t)) ~ ∂x(1.0 / (1.0 / D₀ + exp(α * (c(x, t) - χ)) / D₀) * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 04: ∂t(c(x, t)) ~ ∂x(D₀ / (1.0 + exp(α * (c(x, t) - χ))) * ∂x(c(x, t)))" begin
        diff_eq = ∂t(c(x, t)) ~ ∂x(D₀ / (1.0 + exp(α * (c(x, t) - χ))) * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 05: ∂t(c(x, t)) ~ ∂x(1/x * ∂x(c(x, t)))" begin
        diff_eq = ∂t(c(x, t)) ~ ∂x(1 / x * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 06: ∂t(c(x, t)) ~ ∂x(x*∂x(c(x, t)))/c(x,t)" begin
        diff_eq = ∂t(c(x, t)) ~ ∂x(x * ∂x(c(x, t))) / c(x, t)
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 07: ∂t(c(x, t)) ~ ∂x(1/(1+c(x,t)) ∂x(c(x, t)))" begin
        diff_eq = ∂t(c(x, t)) ~ ∂x(1 / (1 + c(x, t)) * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 08: ∂t(c(x, t)) ~ c(x, t) * ∂x(c(x,t) * ∂x(c(x, t)))" begin
        diff_eq = ∂t(c(x, t)) ~ c(x, t) * ∂x(c(x, t) * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 09: ∂t(c(x, t)) ~ c(x, t) * ∂x(c(x,t) * ∂x(c(x, t)))/(1+c(x,t))" begin
        diff_eq = c(x, t) * ∂x(c(x, t) * ∂x(c(x, t))) / (1 + c(x, t))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 10: ∂t(c(x, t)) ~ c(x, t) * ∂x(c(x,t) * ∂x(c(x, t)))/(1+c(x,t))" begin
        diff_eq = c(x, t) * ∂x(c(x, t)^(-1) * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

    @testset "Test 11: ∂t(c(x, t)) ~ ∂x(1/(1+c(x,t)^2) ∂x(c(x, t)))" begin
        diff_eq = ∂t(c(x, t)) ~ ∂x(1 / (1 + c(x, t)^2) * ∂x(c(x, t)))
        @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)])
        discretization = MOLFiniteDifference([x => Δx], t)
    end

end


using Random
Random.seed!(1)

numbers = 10:25
N = 10^7

firstDigit(x) = Int(floor(x/10))
secondDigit(x) = x%10

numThirteen, numFirstIsOne, numSecondIsThree = 0, 0, 0

for _ in 1:N
    X = rand(numbers)
    global numThirteen += X == 13
    global numFirstIsOne += firstDigit(X) == 1
    global numSecondIsThree += secondDigit(X) == 3
end

probThirteen, probFirstIsOne, probSecondIsThree =
    (numThirteen,numFirstIsOne,numSecondIsThree)./N

println("P(13) = ", round(probThirteen, digits=4),
        "\nP(1_) = ",round(probFirstIsOne, digits=4),
        "\nP(_3) = ", round(probSecondIsThree, digits=4),
        "\nP(1_)*P(_3) = ",round(probFirstIsOne*probSecondIsThree, digits=4))


m = model(data...) # instantiate model on the data
q = vi(m, vi_alg)  # perform VI on `m` using the VI method `vi_alg`, which returns a `VariationalPosterior`

using Random
using Turing
using Turing: Variational

Random.seed!(42);

# generate data
x = randn(2000);

@model model(x) = begin
    s ~ InverseGamma(2, 3)
    m ~ Normal(0.0, sqrt(s))
    for i = 1:length(x)
        x[i] ~ Normal(m, sqrt(s))
    end
end;

# Instantiate model
m = model(x);

samples_nuts = sample(m, NUTS(200, 0.65), 10000);

@doc(Variational.vi)

@doc(Variational.meanfield)

@doc(Variational.ADVI)

# ADVI
advi = ADVI(10, 1000)
q = vi(m, advi);

q isa MultivariateDistribution

rand(q)

logpdf(q, rand(q))

var(x), mean(x)

(mean(rand(q, 1000); dims = 2)..., )

samples = rand(q, 10000);

# setup for plotting
using Plots, LaTeXStrings, StatsPlots

p1 = histogram(samples[1, :], bins=100, normed=true, alpha=0.2, color = :blue, label = "")
density!(samples[1, :], label = "s (ADVI)", color = :blue, linewidth = 2)
density!(samples_nuts, :s; label = "s (NUTS)", color = :green, linewidth = 2)
vline!([var(x)], label = "s (data)", color = :black)
vline!([mean(samples[1, :])], color = :blue, label ="")

p2 = histogram(samples[2, :], bins=100, normed=true, alpha=0.2, color = :blue, label = "")
density!(samples[2, :], label = "m (ADVI)", color = :blue, linewidth = 2)
density!(samples_nuts, :m; label = "m (NUTS)", color = :green, linewidth = 2)
vline!([mean(x)], color = :black, label = "m (data)")
vline!([mean(samples[2, :])], color = :blue, label="")

plot(p1, p2, layout=(2, 1), size=(900, 500))

# used to compute closed form expression of posterior
using ConjugatePriors

# closed form computation
# notation mapping has been verified by explicitly computing expressions
# in "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
μ₀ = 0.0 # => μ
κ₀ = 1.0 # => ν, which scales the precision of the Normal
α₀ = 2.0 # => "shape"
β₀ = 3.0 # => "rate", which is 1 / θ, where θ is "scale"

# prior
pri = NormalGamma(μ₀, κ₀, α₀, β₀)

# posterior
post = posterior(pri, Normal, x)

# marginal distribution of τ = 1 / σ²
# Eq. (90) in "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
# `scale(post)` = θ
p_τ = Gamma(post.shape, scale(post))
p_σ²_pdf = z -> pdf(p_τ, 1 / z) # τ => 1 / σ²

# marginal of μ
# Eq. (91) in "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
p_μ = TDist(2 * post.shape)

μₙ = post.mu    # μ → μ
κₙ = post.nu    # κ → ν
αₙ = post.shape # α → shape
βₙ = post.rate  # β → rate

# numerically more stable but doesn't seem to have effect; issue is probably internal to
# `pdf` which needs to compute ≈ Γ(1000)
p_μ_pdf = z -> exp(logpdf(p_μ, (z - μₙ) * exp(- 0.5 * log(βₙ) + 0.5 * log(αₙ) + 0.5 * log(κₙ))))

# posterior plots
p1 = plot();
histogram!(samples[1, :], bins=100, normed=true, alpha=0.2, color = :blue, label = "")
density!(samples[1, :], label = "s (ADVI)", color = :blue)
density!(samples_nuts, :s; label = "s (NUTS)", color = :green)
vline!([mean(samples[1, :])], linewidth = 1.5, color = :blue, label ="")

# normalize using Riemann approx. because of (almost certainly) numerical issues
Δ = 0.001
r = 0.75:0.001:1.50
norm_const = sum(p_σ²_pdf.(r) .* Δ)
plot!(r, p_σ²_pdf, label = "s (posterior)", color = :red);
vline!([var(x)], label = "s (data)", linewidth = 1.5, color = :black, alpha = 0.7);
xlims!(0.75, 1.35);

p2 = plot();
histogram!(samples[2, :], bins=100, normed=true, alpha=0.2, color = :blue, label = "")
density!(samples[2, :], label = "m (ADVI)", color = :blue)
density!(samples_nuts, :m; label = "m (NUTS)", color = :green)
vline!([mean(samples[2, :])], linewidth = 1.5, color = :blue, label="")

# normalize using Riemann approx. because of (almost certainly) numerical issues
Δ = 0.0001
r = -0.1 + mean(x):Δ:0.1 + mean(x)
norm_const = sum(p_μ_pdf.(r) .* Δ)
plot!(r, z -> p_μ_pdf(z) / norm_const, label = "m (posterior)", color = :red);
vline!([mean(x)], label = "m (data)", linewidth = 1.5, color = :black, alpha = 0.7);

xlims!(-0.25, 0.25);

p = plot(p1, p2; layout=(2, 1), size=(900, 500))

Random.seed!(1);

# Import RDatasets.
using RDatasets

# Hide the progress prompt while sampling.
Turing.setprogress!(false);

# Import the "Default" dataset.
data = RDatasets.dataset("datasets", "mtcars");

# Show the first six rows of the dataset.
first(data, 6)

# Function to split samples.
function split_data(df, at = 0.70)
    r = size(df,1)
    index = Int(round(r * at))
    train = df[1:index, :]
    test  = df[(index+1):end, :]
    return train, test
end

# A handy helper function to rescale our dataset.
function standardize(x)
    return (x .- mean(x, dims=1)) ./ std(x, dims=1), x
end

# Another helper function to unstandardize our datasets.
function unstandardize(x, orig)
    return (x .+ mean(orig, dims=1)) .* std(orig, dims=1)
end

# Remove the model column.
select!(data, Not(:Model))

# Standardize our dataset.
(std_data, data_arr) = standardize(Matrix(data))

# Split our dataset 70%/30% into training/test sets.
train, test = split_data(std_data, 0.7)

# Save dataframe versions of our dataset.
train_cut = DataFrame(train, names(data))
test_cut = DataFrame(test, names(data))

# Create our labels. These are the values we are trying to predict.
train_label = train_cut[:, :MPG]
test_label = test_cut[:, :MPG]

# Get the list of columns to keep.
remove_names = filter(x->!in(x, [:MPG, :Model]), names(data))

# Filter the test and train sets.
train = Matrix(train_cut[:,remove_names]);
test = Matrix(test_cut[:,remove_names]);

# Bayesian linear regression.
@model linear_regression(x, y, n_obs, n_vars, ::Type{T}=Vector{Float64}) where {T} = begin
    # Set variance prior.
    σ₂ ~ truncated(Normal(0,100), 0, Inf)

    # Set intercept prior.
    intercept ~ Normal(0, 3)

    # Set the priors on our coefficients.
    coefficients ~ MvNormal(zeros(n_vars), 10 * ones(n_vars))

    # Calculate all the mu terms.
    mu = intercept .+ x * coefficients
    y ~ MvNormal(mu, σ₂)
end;

n_obs, n_vars = size(train)
m = linear_regression(train, train_label, n_obs, n_vars);

q0 = Variational.meanfield(m)
typeof(q0)

advi = ADVI(10, 10_000)

using Flux, Turing
using Turing.Variational

vi(m, advi; optimizer = Flux.ADAM())

opt = Variational.DecayedADAGrad(1e-2, 1.1, 0.9)

q = vi(m, advi, q0; optimizer = opt)
typeof(q)

z = rand(q, 10_000);

avg = vec(mean(z; dims = 2))

_, sym2range = bijector(m, Val(true));
sym2range

avg[union(sym2range[:σ₂]...)]

avg[union(sym2range[:intercept]...)]

avg[union(sym2range[:coefficients]...)]

function plot_variational_marginals(z, sym2range)
    ps = []

    for (i, sym) in enumerate(keys(sym2range))
        indices = union(sym2range[sym]...)  # <= array of ranges
        if sum(length.(indices)) > 1
            offset = 1
            for r in indices
                for j in r
                    p = density(z[j, :], title = "$(sym)[$offset]", titlefontsize = 10, label = "")
                    push!(ps, p)

                    offset += 1
                end
            end
        else
            p = density(z[first(indices), :], title = "$(sym)", titlefontsize = 10, label = "")
            push!(ps, p)
        end
    end

    return plot(ps..., layout = (length(ps), 1), size = (500, 1500))
end

plot_variational_marginals(z, sym2range)

chain = sample(m, NUTS(0.65), 10_000);

plot(chain)

vi_mean = vec(mean(z; dims = 2))[[union(sym2range[:coefficients]...)..., union(sym2range[:intercept]...)..., union(sym2range[:σ₂]...)...]]

mean(chain).nt.mean

sum(abs2, mean(chain).nt.mean .- vi_mean)

# Import the GLM package.
using GLM

# Perform multivariate OLS.
ols = lm(@formula(MPG ~ Cyl + Disp + HP + DRat + WT + QSec + VS + AM + Gear + Carb), train_cut)

# Store our predictions in the original dataframe.
train_cut.OLSPrediction = unstandardize(GLM.predict(ols), data.MPG);
test_cut.OLSPrediction = unstandardize(GLM.predict(ols, test_cut), data.MPG);

# Make a prediction given an input vector.
function prediction_chain(chain, x)
    p = get_params(chain)
    α = mean(p.intercept)
    β = collect(mean.(p.coefficients))
    return  α .+ x * β
end

# Make a prediction using samples from the variational posterior given an input vector.
function prediction(samples::AbstractVector, sym2ranges, x)
    α = mean(samples[union(sym2ranges[:intercept]...)])
    β = vec(mean(samples[union(sym2ranges[:coefficients]...)]; dims = 2))
    return  α .+ x * β
end

function prediction(samples::AbstractMatrix, sym2ranges, x)
    α = mean(samples[union(sym2ranges[:intercept]...), :])
    β = vec(mean(samples[union(sym2ranges[:coefficients]...), :]; dims = 2))
    return  α .+ x * β
end

# Unstandardize the dependent variable.
train_cut.MPG = unstandardize(train_cut.MPG, data.MPG);
test_cut.MPG = unstandardize(test_cut.MPG, data.MPG);

# Show the first side rows of the modified dataframe.
first(test_cut, 6)

z = rand(q, 10_000);

# Calculate the predictions for the training and testing sets using the samples `z` from variational posterior
train_cut.VIPredictions = unstandardize(prediction(z, sym2range, train), data.MPG);
test_cut.VIPredictions = unstandardize(prediction(z, sym2range, test), data.MPG);

train_cut.BayesPredictions = unstandardize(prediction_chain(chain, train), data.MPG);
test_cut.BayesPredictions = unstandardize(prediction_chain(chain, test), data.MPG);

vi_loss1 = mean((train_cut.VIPredictions - train_cut.MPG).^2)
bayes_loss1 = mean((train_cut.BayesPredictions - train_cut.MPG).^2)
ols_loss1 = mean((train_cut.OLSPrediction - train_cut.MPG).^2)

vi_loss2 = mean((test_cut.VIPredictions - test_cut.MPG).^2)
bayes_loss2 = mean((test_cut.BayesPredictions - test_cut.MPG).^2)
ols_loss2 = mean((test_cut.OLSPrediction - test_cut.MPG).^2)

println("Training set:
    VI loss: $vi_loss1
    Bayes loss: $bayes_loss1
    OLS loss: $ols_loss1
Test set:
    VI loss: $vi_loss2
    Bayes loss: $bayes_loss2
    OLS loss: $ols_loss2")

z = rand(q, 1000);
preds = hcat([unstandardize(prediction(z[:, i], sym2range, test), data.MPG) for i = 1:size(z, 2)]...);

scatter(1:size(test, 1), mean(preds; dims = 2), yerr=std(preds; dims = 2), label="prediction (mean ± std)", size = (900, 500), markersize = 8)
scatter!(1:size(test, 1), unstandardize(test_label, data.MPG), label="true")
xaxis!(1:size(test, 1))
ylims!(95, 140)
title!("Mean-field ADVI (Normal)")

preds = hcat([unstandardize(prediction_chain(chain[i], test), data.MPG) for i = 1:5:size(chain, 1)]...);

scatter(1:size(test, 1), mean(preds; dims = 2), yerr=std(preds; dims = 2), label="prediction (mean ± std)", size = (900, 500), markersize = 8)
scatter!(1:size(test, 1), unstandardize(test_label, data.MPG), label="true")
xaxis!(1:size(test, 1))
ylims!(95, 140)
title!("MCMC (NUTS)")

using Bijectors

using Bijectors: Scale, Shift

d = length(q)
base_dist = Turing.DistributionsAD.TuringDiagMvNormal(zeros(d), ones(d))

to_constrained = inv(bijector(m));

function getq(θ)
    d = length(θ) ÷ 2
    A = @inbounds θ[1:d]
    b = @inbounds θ[d + 1: 2 * d]

    b = to_constrained ∘ Shift(b; dim = Val(1)) ∘ Scale(exp.(A); dim = Val(1))

    return transformed(base_dist, b)
end

q_mf_normal = vi(m, advi, getq, randn(2 * d));

p1 = plot_variational_marginals(rand(q_mf_normal, 10_000), sym2range) # MvDiagNormal + Affine transformation + to_constrained
p2 = plot_variational_marginals(rand(q, 10_000), sym2range)  # Turing.meanfield(m)

plot(p1, p2, layout = (1, 2), size = (800, 2000))

using LinearAlgebra

# Using `ComponentArrays.jl` together with `UnPack.jl` makes our lives much easier.
using ComponentArrays, UnPack

proto_arr = ComponentArray(
    L = zeros(d, d),
    b = zeros(d)
)
proto_axes = proto_arr |> getaxes
num_params = length(proto_arr)

function getq(θ)
    L, b = begin
        @unpack L, b = ComponentArray(θ, proto_axes)
        LowerTriangular(L), b
    end
    # For this to represent a covariance matrix we need to ensure that the diagonal is positive.
    # We can enforce this by zeroing out the diagonal and then adding back the diagonal exponentiated.
    D = Diagonal(diag(L))
    A = L - D + exp(D) # exp for Diagonal is the same as exponentiating only the diagonal entries

    b = to_constrained ∘ Shift(b; dim = Val(1)) ∘ Scale(A; dim = Val(1))

    return transformed(base_dist, b)
end

advi = ADVI(10, 20_000)

q_full_normal = vi(m, advi, getq, randn(num_params); optimizer = Variational.DecayedADAGrad(1e-2));

A = q_full_normal.transform.ts[1].a

heatmap(cov(A * A'))

zs = rand(q_full_normal, 10_000);

p1 = plot_variational_marginals(rand(q_mf_normal, 10_000), sym2range)
p2 = plot_variational_marginals(rand(q_full_normal, 10_000), sym2range)

plot(p1, p2, layout = (1, 2), size = (800, 2000))

# Unfortunately, it seems like this has quite a high variance which is likely to be due to numerical instability,
# so we consider a larger number of samples. If we get a couple of outliers due to numerical issues,
# these kind affect the mean prediction greatly.
z = rand(q_full_normal, 10_000);

train_cut.VIFullPredictions = unstandardize(prediction(z, sym2range, train), data.MPG);
test_cut.VIFullPredictions = unstandardize(prediction(z, sym2range, test), data.MPG);

vi_loss1 = mean((train_cut.VIPredictions - train_cut.MPG).^2)
vifull_loss1 = mean((train_cut.VIFullPredictions - train_cut.MPG).^2)
bayes_loss1 = mean((train_cut.BayesPredictions - train_cut.MPG).^2)
ols_loss1 = mean((train_cut.OLSPrediction - train_cut.MPG).^2)

vi_loss2 = mean((test_cut.VIPredictions - test_cut.MPG).^2)
vifull_loss2 = mean((test_cut.VIFullPredictions - test_cut.MPG).^2)
bayes_loss2 = mean((test_cut.BayesPredictions - test_cut.MPG).^2)
ols_loss2 = mean((test_cut.OLSPrediction - test_cut.MPG).^2)

println("Training set:
    VI loss: $vi_loss1
    Bayes loss: $bayes_loss1
    OLS loss: $ols_loss1
Test set:
    VI loss: $vi_loss2
    Bayes loss: $bayes_loss2
    OLS loss: $ols_loss2")

z = rand(q_mf_normal, 1000);
preds = hcat([unstandardize(prediction(z[:, i], sym2range, test), data.MPG) for i = 1:size(z, 2)]...);

p1 = scatter(1:size(test, 1), mean(preds; dims = 2), yerr=std(preds; dims = 2), label="prediction (mean ± std)", size = (900, 500), markersize = 8)
scatter!(1:size(test, 1), unstandardize(test_label, data.MPG), label="true")
xaxis!(1:size(test, 1))
ylims!(95, 140)
title!("Mean-field ADVI (Normal)")

z = rand(q_full_normal, 1000);
preds = hcat([unstandardize(prediction(z[:, i], sym2range, test), data.MPG) for i = 1:size(z, 2)]...);

p2 = scatter(1:size(test, 1), mean(preds; dims = 2), yerr=std(preds; dims = 2), label="prediction (mean ± std)", size = (900, 500), markersize = 8)
scatter!(1:size(test, 1), unstandardize(test_label, data.MPG), label="true")
xaxis!(1:size(test, 1))
ylims!(95, 140)
title!("Full ADVI (Normal)")

preds = hcat([unstandardize(prediction_chain(chain[i], test), data.MPG) for i = 1:5:size(chain, 1)]...);

p3 = scatter(1:size(test, 1), mean(preds; dims = 2), yerr=std(preds; dims = 2), label="prediction (mean ± std)", size = (900, 500), markersize = 8)
scatter!(1:size(test, 1), unstandardize(test_label, data.MPG), label="true")
xaxis!(1:size(test, 1))
ylims!(95, 140)
title!("MCMC (NUTS)")

plot(p1, p2, p3, layout = (1, 3), size = (900, 250), label="")

if isdefined(Main, :TuringTutorials)
    Main.TuringTutorials.tutorial_footer(WEAVE_ARGS[:folder], WEAVE_ARGS[:file])
end


## Preloads
module DHC_Utils

    using Statistics
    using FFTW
    using Plots
    using LinearAlgebra
    using StaticArrays
    using HybridArrays

    export finklet
    export fink_filter_bank_slow
    export fink_filter_bank
    export fink_filter_list
    export speedy_DHC

    ## Define wavelet in Fourier space
    function finklet(j, l)
        # -------- set filters
        jrad = 7-j
        dθ = π/8        # 8 angular bins hardwired
        θ_l = dθ*l
        # -------- define coordinates
        nx = 256
        xbox = LinRange(-nx/2, nx/2-1 , nx)
        # make a 256x256 grid of X
        sx = xbox' .* ones(nx)
        sy = ones(nx)' .* xbox
        r  = sqrt.((sx).^2 + (sy).^2)
        θ  = mod.(atan.(sy, sx).+π .-θ_l,2*π)
        nozeros = r .> 0
        logr = log2.(r[nozeros])
        r[nozeros] = logr
        # -------- in Fourier plane, envelope of psi_j,l
        mask = (abs.(θ.-π).<= dθ) .& (abs.(r.-jrad) .<= 1)
        # -------- angular part
        ang = cos.((θ.-π).*4)
        # -------- radial part
        rad = cos.((r.-jrad).*π./2)
        psi = mask.*ang.*rad             #mask times angular part times radial part
        return psi
    end

    ## Compute the whole filter bank.  Legacy code for comparison
    function fink_filter_bank_slow(J,L)
        fink_filter = Array{Float64, 4}(undef, 256, 256, J, L)
        for l = 1:L
            for j = 1:J
                @inbounds fink_filter[:,:,j,l]=fftshift(finklet(j-1,l-1))
            end
        end
        return fink_filter
    end

    ## Faster filter bank generation.  Only square filters allowed.
    function fink_filter_bank(J::Integer, L::Integer; nx::Integer=256, wid::Integer=1)

        # -------- set parameters
        dθ   = π/8        # 8 angular bins hardwired
        dx   = nx÷2-1

        # -------- allocate output array of zeros
        filt = zeros(Float64, nx, nx, J, L)

        # -------- allocate theta and logr arrays
        logr = zeros(Float64, nx, nx)
        θ    = zeros(Float64, nx, nx)

        for l = 0:L-1
            θ_l = dθ*l

        # -------- allocate anggood BitArray
            anggood = falses(nx, nx)

        # -------- loop over pixels
            for x = 1:nx
                sx = mod(x+dx, nx)-dx -1    # define sx,sy so that no fftshift() needed
                for y = 1:nx
                    sy = mod(y+dx, nx)-dx -1
                    θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                    θ_good = abs(θ_pix-π) <= (dθ*wid)

                    # If this is a pixel we might use, calculate log2(r)
                    if θ_good
                        anggood[y, x] = θ_good
                        θ[y, x]       = θ_pix
                        r = sqrt(sx^2 + sy^2)
                        logr[y, x] = log2(max(1,r))
                    end
                end
            end
            angmask = findall(anggood)
        # -------- compute the wavelet in the Fourier domain
        #          the angular factor is the same for all j
            F_angular = cos.((θ[angmask].-π).*(4/wid))

        # -------- loop over j for the radial part
            for j = 0:J-1
                jrad  = 7-j
                Δj    = abs.(logr[angmask].-jrad)
                rmask = (Δj .<= 1)

        # -------- radial part
                F_radial = cos.(Δj[rmask] .* (π/2))
                ind      = angmask[rmask]
                filt[ind,j+1,l+1] = F_radial .* F_angular[rmask]
            end
        end
        return filt
    end

    ## Make a list of non-zero pixels for the (Fourier plane) filters
    function fink_filter_list(filt)
        (ny,nx,J,L) = size(filt)

        # Allocate output arrays
        filtind = fill(CartesianIndex{2}[], J, L)
        filtval = fill(Float64[], J, L)

        # Loop over J,L and record non-zero values
        for l=1:L
            for j=1:J
                f = @view filt[:,:,j,l]
                ind = findall(f .> 1E-13)
                val = f[ind]
                filtind[j,l] = ind
                filtval[j,l] = val
            end
        end
        return [filtind, filtval]
    end

    ## Todo list
    # Check if 2 threads really help FFT when computer is busy
    function speedy_DHC(image::Array{Float64,2}, filter_list)
        # Use 2 threads for FFT
        FFTW.set_num_threads(2)

        # array sizes
        (Nx, Ny)  = size(image)
        (J,L)     = size(filter_list[1])

        # allocate coeff arrays
        out_coeff = []
        S0  = zeros(Float64, 2)
        S1  = zeros(Float64, J, L)
        S20 = zeros(Float64, J, L, J, L)
        S12 = zeros(Float64, J, L, J, L)

        # allocate image arrays for internal use
        im_fdf_0_1 = zeros(Float64,           Nx, Ny, J, L)   # this must be zeroed!
        im_rd_0_1  = Array{Float64, 4}(undef, Nx, Ny, J, L)

        ## 0th Order
        S0[1]   = mean(image)
        norm_im = image.-S0[1]
        S0[2]   = sum(norm_im .* norm_im)/(Nx*Ny)
        norm_im ./= sqrt(Nx*Ny*S0[2])

        append!(out_coeff,S0[:])

        ## 1st Order
        im_fd_0 = fft(norm_im)

        # unpack filter_list
        f_ind   = filter_list[1]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_list[2]

        zarr = zeros(ComplexF64, Nx, Ny)  # temporary array to fill with zvals

        # make a FFTW "plan" for an array of the given size and type
        P   = plan_ifft(im_fd_0)   # P is an operator, P*im is ifft(im)

        ## Main 1st Order and Precompute 2nd Order
        for l = 1:L
            for j = 1:J
                S1tot = 0.0
                f_i = f_ind[j,l]  # CartesianIndex list for filter
                f_v = f_val[j,l]  # Values for f_i
                # for (ind, val) in zip(f_i, f_v)   # this is slower!
                if length(f_i) > 0
                    for i = 1:length(f_i)
                        ind       = f_i[i]
                        zval      = f_v[i] * im_fd_0[ind]
                        S1tot    += abs2(zval)
                        zarr[ind] = zval
                        im_fdf_0_1[ind,j,l] = abs(zval)
                    end
                    S1[j,l] = S1tot
                    im_rd_0_1[:,:,j,l] .= abs2.(P*zarr)
                    zarr[f_i] .= 0
                end
            end
        end
        append!(out_coeff, S1[:])

        # we stored the abs()^2, so take sqrt (this is faster to do all at once)
        im_rd_0_1 .= sqrt.(im_rd_0_1)

        ## 2nd Order
        Amat = reshape(im_fdf_0_1, Nx*Ny, J*L)
        S12  = reshape(Amat' * Amat, J, L, J, L)
        Amat = reshape(im_rd_0_1, Nx*Ny, J*L)
        S20  = reshape(Amat' * Amat, J, L, J, L)

        append!(out_coeff, S20)
        append!(out_coeff, S12)

        return out_coeff
    end

end # of module


# The cube, 41063625 (345^3), can be permuted to produce two other cubes:
# 56623104 (384^33) and 66430125 (405^33). In fact, 41063625 is the smallest
# cube which has exactly three permutations of its digits which are also cube.
#
# Find the smallest cube for which exactly five permutations of its digits are
# cube.

using ProjectEulerSolutions

# For each fixed number of digits, search through the right ranges of numbers
# to build up sets of cubes.  Store them as sorted tuples, counting the number
# that match to the same key.  When we hit 5, return the smallest.
function p062solution(n::Integer=5)::Integer

    ndig_max = ndigits(2^63-1) - 1  # Bigger than 18 and we will need to use BigInts
    for ndig in 9:ndig_max

        dict = Dict{Tuple, Array{Integer,1}}()

        num_low = Integer(ceil(10^((ndig-1)/3)))
        num_high = Integer(floor(10^(ndig/3)))
        max_count = 1
        max_key = 1

        for i in num_low:num_high
            num = i^3
            key = Tuple(sort(digits(num)))
            if haskey(dict, key)
                append!(dict[key], num)
            else
                dict[key] = [num]
            end
            c = length(dict[key])
            if max_count < c
               max_count = c
               if max_count == n
                   return minimum(dict[key])
               end
            end
        end
    end
    return -1
end

p062 = Problems.Problem(p062solution)

Problems.benchmark(p062, 5)


#=
Given a list of integers, write a function that returns the largest sum of non-adjacent numbers. Numbers can be 0 or negative.

For example, [2, 4, 6, 2, 5] should return 13, since we pick 2, 6, and 5. [5, 1, 1, 5] should return 10, since we pick 5 and 5.

Follow-up: Can you do this in O(N) time and constant space?
=#

using Test
include("Solutions/problem9_sum_notadjacentnumbers.jl")
@test sum_notadjacent([2, 4, 6, 2, 5]) == 13
@test sum_notadjacent([5, 1, 1, 5]) == 10
@test sum_notadjacent([-3, -4, -7, 1, 0, 1]) == 2


function rand_candecomp{N}(r::Int64, dims::NTuple{N, Int};
                           lambdas_nonneg::Bool=false, factors_nonneg::Bool=false)
    rnd_factor = factors_nonneg ? x -> abs.(randn(x...)) : randn
    rnd_lambda = lambdas_nonneg ? x -> abs.(randn(x...)) : randn
    return CANDECOMP((Matrix{Float64}[rnd_factor((s, r)) for s in dims]...),
                     rnd_lambda(r))
end

rand_kruskal3{N}(r::Int64, dims::NTuple{N, Int}, nonnegative::Bool) =
    compose(rand_candecomp(r, dims, lambdas_nonneg=nonnegative, factors_nonneg=nonnegative))

function rand_tucker{N}(core_dims::NTuple{N, Int}, dims::NTuple{N, Int};
                        core_nonneg::Bool=false, factors_nonneg::Bool=false)
    rnd_factor = factors_nonneg ? x -> abs.(randn(x...)) : randn
    rnd_core = core_nonneg ? x -> abs.(randn(x...)) : randn
    return Tucker((Matrix{Float64}[rnd_factor((dims[i], core_dims[i])) for i in 1:N]...),
                  rnd_core(core_dims))
end

function add_noise{T, N}(tnsr::Array{T,N}, sn_ratio = 0.6, nonnegative::Bool = false)
    tnsr_noise = randn(size(tnsr)...)
    if nonnegative
        map!(x -> max(0.0, x), tnsr_noise, tnsr_noise)
    end
    tnsr + 10^(-sn_ratio/0.2)*vecnorm(tnsr)/vecnorm(tnsr)*tnsr_noise
end


using Random, Distributions, StatsPlots, Plots; pyplot()
Random.seed!(0)

mu = 20
d1, d2 = Normal(mu,mu), Exponential(mu)

n = 100
data1 = rand(d1,n)
data2 = rand(d2,n)

qqnorm(data1, c=:blue, ms=3, msw=0, label="Normal Data")
qqnorm!(data2, c=:red, ms=3, msw=0, label="Exponential Data",
        xlabel="Normal Theoretical Quantiles",
        ylabel="Data Quantiles", legend=true)


#
# Implements example 4.2 in Sarkka book, but due to technical reasons with
# regard to handling univariate and multivariate gaussians consistently,
# we simulate two random walks by making A, H, Q, R matrices as opposed to
# scalars specified in book.
#

function Example_4_2(; T = 400, seed = 1)

    # Set random number generator

    rg = MersenneTwister(seed)

    # Matrices in the dynamic model are:

    A = H = diagm(ones(2))
    Q = R = diagm(ones(2))

    # initialise

    m₀, P₀ = zeros(2), diagm(ones(2))

    x₀ = rand(rg, MvNormal(m₀, P₀))

    # store measurements and states

    yclean = Array{Vector{Float64},1}(undef, T)

    y = Array{Vector{Float64},1}(undef, T)

    x = Array{Vector{Float64},1}(undef, T)

    # first state

    x[1]      = A * x₀

    y[1]      = rand(rg, MvNormal(H * x[1], R))

    yclean[1] = H * x₀

    for k in 2:T

        x[k] = rand(rg, MvNormal(A * x[k-1], Q)) # (4.32)

        yclean[k] = H * x[k]

        y[k] = rand(rg, MvNormal(H * (x[k]), R)) # (4.33)

    end

    return yclean, y, x, A, H, Q, R, m₀, P₀

end

function runExample_4_2(seed = 1)

    # Simulate data

    yclean, y, x, A, H, Q, R, m₀, P₀ = Example_4_2(seed=seed)

    # Inference

    μfilter, Σfilter = filteringrecursion(y; A = A, H = H, Q = Q, R = R, m₀ = m₀, P₀ = P₀)

    μsmooth, Σsmooth = smoothingrecursion(y; A = A, H = H, Q = Q, R = R, m = μfilter, P = Σfilter)

    # collect positions

    truepositions = reduce(hcat, [xₖ[1:2] for xₖ in x])

    filtpositions = reduce(hcat, [xₖ[1:2] for xₖ in μfilter])

    smoopositions = reduce(hcat, [xₖ[1:2] for xₖ in μsmooth])

    # plot positions

    figure(1); cla()

    plot(truepositions[1,:], "-", label = "true")

    plot(filtpositions[1,:], "-", label = "filter")

    plot(smoopositions[1,:], "-", label = "smoother")

    legend()

    nothing

end


# 'cstrt_tamvakis_ngon' provides the vertices coordinates (a,b) of the
# Tamvakis n-gon
# Please, cite...
#   N. K. Tamvakis. On the perimeter and the area of the convex polygon of
#   a given diameter. Bull. Greek Math. Soc, 28: 115-132, 1987.
function cstrt_tamvakis_ngon(n)
if n>=3
    a = zeros(n-1); b = zeros(n-1);
    if mod(n,3) == 0
        a[Int(n/3)] = 1/2; a[Int(2*n/3)] = -a[Int(n/3)];
        b[Int(n/3)] = sqrt(3)/2; b[Int(2*n/3)] = b[Int(n/3)];
        if n>3
            for j = 1:Int((n-3)/3)
                a[Int(n/3+j)] = cos(pi/3+j*pi/n);
                b[Int(n/3+j)] = sin(pi/3+j*pi/n);
                a[Int(2*n/3+j)] = a[Int(n/3)] + cos(pi+j*pi/n);
                b[Int(2*n/3+j)] = b[Int(n/3)] + sin(pi+j*pi/n);
                a[Int(n/3-j)] = -a[Int(2*n/3+j)];
                b[Int(n/3-j)] = b[Int(2*n/3+j)];
            end
        end
    elseif mod(n,3) == 1
        a[Int((n-1)/3)] = 1/2; a[Int((2*n+1)/3)] = -a[Int((n-1)/3)];
        b[Int((n-1)/3)] = sqrt(3)/2; b[Int((2*n+1)/3)] = b[Int((n-1)/3)];
        for j = 1:Int((n-1)/3)
            a[Int((n-1)/3+j)] = cos(pi/3+j*pi/(n+2));
            b[Int((n-1)/3+j)] = sin(pi/3+j*pi/(n+2));
        end
        if n>4
            for j = 1:Int((n-4)/3)
                a[Int((2*n+1)/3+j)] = a[Int((n-1)/3)] + cos(pi+j*pi/(n-1));
                b[Int((2*n+1)/3+j)] = b[Int((n-1)/3)] + sin(pi+j*pi/(n-1));
                a[Int((n-1)/3-j)] = -a[Int((2*n+1)/3+j)];
                b[Int((n-1)/3-j)] = b[Int((2*n+1)/3+j)];
            end
        end
    else
        a[Int((n+1)/3)] = 1/2; a[Int((2*n-1)/3)] = -a[Int((n+1)/3)];
        b[Int((n+1)/3)] = sqrt(3)/2; b[Int((2*n-1)/3)] = b[Int((n+1)/3)];
        for j = 1:Int((n-2)/3)
            a[Int((2*n-1)/3+j)] = a[Int((n+1)/3)] + cos(pi+j*pi/(n+1));
            b[Int((2*n-1)/3+j)] = b[Int((n+1)/3)] + sin(pi+j*pi/(n+1));
            a[Int((n+1)/3-j)] = -a[Int((2*n-1)/3+j)];
            b[Int((n+1)/3-j)] = b[Int((2*n-1)/3+j)];
        end
        if n>5
            for j = 1:Int((n-5)/3)
                a[Int((n+1)/3+j)] = cos(pi/3+j*pi/(n-2));
                b[Int((n+1)/3+j)] = sin(pi/3+j*pi/(n-2));
            end
        end
    end
    return(a,b)
end
end


using DataDrivenDiffEq
using ModelingToolkit
using OrdinaryDiffEq
using LinearAlgebra
using Plots
gr()

# Create a test problem
function pendulum(u, p, t)
    x = u[2]
    y = -9.81sin(u[1]) - 0.1u[2]
    return [x;y]
end

u0 = [0.2π; -1.0]
tspan = (0.0, 40.0)
prob = ODEProblem(pendulum, u0, tspan)
sol = solve(prob)

plot(sol)

# Create the differential data
DX = similar(sol[:,:])
for (i, xi) in enumerate(eachcol(sol[:,:]))
    DX[:,i] = pendulum(xi, [], 0.0)
end

# Create a basis
@variables u[1:2]

# Lots of polynomials
polys = [u[1]^0]
for i ∈ 1:3
    for j ∈ 1:3
        push!(polys, u[1]^i*u[2]^j)
    end
end

# And some other stuff
h = [1u[1];1u[2]; cos(u[1]); sin(u[1]); u[1]*u[2]; u[1]*sin(u[2]); u[2]*cos(u[2]); polys...]

basis = Basis(h, u)

# Get the reduced basis via the sparse regression
Ψ = SInDy(sol[:,:], DX, basis, ϵ = 1e-10)

# Transform into ODE System
sys = ODESystem(Ψ)

# Simulate
estimator = ODEProblem(dynamics(Ψ), u0, tspan)
sol_ = solve(estimator, saveat = sol.t)

# Yeah! We got it right
plot(sol, vars = (1,2))
plot!(sol, vars = (1,2))

norm(sol-sol_) # ≈ 1.89e-13


import FastTransforms
import QuadratureRules: chebyshev_points, shift_nodes

@testset "$(rpad("Gauß-Chebyshev",80))" begin

    for s in 2:10
        @test GaussChebyshevQuadrature(s) == ChebyshevQuadrature(s, 1)
        @test GaussChebyshevQuadrature(s) == ChebyshevQuadrature(Float64, s, Val(1))

        @test sum(weights(GaussChebyshevQuadrature(s))) ≈ 1

        c = shift_nodes(reverse(FastTransforms.chebyshevpoints(Float64, s, Val(1))))

        @test chebyshev_points(Float64, s, Val(1)) ≈ c
        @test nodes(GaussChebyshevQuadrature(s))   ≈ c

        # TODO test weights

    end

end


using Pkg
Pkg.activate(pwd())

# # Wave equation
# ## Solving the wave equation

struct Wave
    f
    g
    c
end

# ### Exercise:
#
# Write the function `solve_wave(T, L, wave::Wave; n_t=100, n_x=100)` that solves the wave
# equation.
#
# **Hint**: Follow the procedure from the previous exercise. Discretize time and space,
# initialize the solution, add the boundary conditions, add the initial conditions and
# finally, iterate over time.
#
# ---
# ### Solution:

# ---

using Plots

function plot_wave(y, file_name; fps = 60, kwargs...)
    anim = @animate for (i, y_row) in enumerate(eachrow(y))
        plot(
            y_row;
            title = "t = $(i-1)Δt",
            xlabel = "x",
            ylabel = "y(t, x)",
            legend = false,
            linewidth = 2,
            kwargs...
        )
    end
    gif(anim, file_name; fps, show_msg = false)

    return nothing
end

# ### Exercise:
#
# Solve the wave equation for $L=\frac32\pi$, $T=240$, $c=0.02$ and the initial
# conditions
#
# $$
# \begin{aligned}
# f(x) &= 2e^{-(x-\frac L2)^2} + \frac{x}{L}, \\
# g(x) &= 0.
# \end{aligned}
# $$
#
# Use time discretization with stepsize $\Delta t=1$ and the space discretization with
# number of points $n_x=101$ and $n_x=7$ steps. Plot two graphs.
#
# ---
# ### Solution:


#=
							Multiples of 3 and 5

	If we list all the natural numbers below 10 that are multiples of 3 or 5, we
	get 3, 5, 6 and 9. The sum of these multiples is 23.

	Find the sum of all the multiples of 3 or 5 below 1000.

							Answer: 233168
=#

function multiples(a::Int, b::Int, limit::Int)
    array = []

    for i in 2:limit-1
        if 0 == i%a || 0 == i%b
            push!(array, i)
        end
    end

    return array
end

println(sum(multiples(3, 5, 1000)))


using JLD
using PyPlot
using StatsBase
using Distributions
using VBMatrixFactorization

function toy_matrix(L, M, H, std)
    # B is a random matrix whose columns we select from
    B = randn(L, H)
    # A is the sparse matrix that selects the columns from A
    A = zeros(M, H)
    for m in 1:M
        A[m, rand(1:H)] = 1
    end

    Y = B*A' + rand(Normal(0, std), L, M)
    return Y, A, B
end

L = 10
M = 20
H = 2
Y_toy, A_toy, B_toy = toy_matrix(L, M, H, 0.05);

# first, initialize a variable of the vbmf_parameters type
# then use it to run the basic vbmf algorithm
# here, the whole run is saved into the data_path folder
# if desc is not specified, then the subfolder with the actual data
# is named using date and time

println(" ----------- Basic VB Matrix factorization ---------------- ")
println("")
params_init = VBMatrixFactorization.vbmf_init(Y_toy, H, ca = 0.1, cb = 0.1, sigma2 = 0.1);
data_path = "./data"
res_vbmf = VBMatrixFactorization.vbmf(Y_toy, params_init, 100, est_covs = true, est_var = true, verb = true,
    logdir = data_path, desc = "vbmf_test");
err = norm(Y_toy - res_vbmf.YHat)
println("||Y - Yhat|| = $err")

# this is how the log is loaded and used
vbmf_log, data, priors = VBMatrixFactorization.load_log(string(data_path,"/vbmf_test"));
params_vbmf = VBMatrixFactorization.vbmf_parameters(); # dummy variable to store individual step data

iter = 3
time_slice = VBMatrixFactorization.extract_params!(vbmf_log, iter, params_vbmf);
err = norm(Y_toy - time_slice.YHat)
println("in iteration number $iter, the error was ||Y - Yhat|| = $err")
println("")

# vbmf_sparse usage
println(" ----------- VB Matrix factorization with sparse A ---------------- ")
println("")
params_sparse_init = VBMatrixFactorization.vbmf_sparse_init(Y_toy, H, ca = 0.1, cb = 0.1, sigma = 0.1);
params_sparse_init.AHat = A_toy
sparse_vbmf = VBMatrixFactorization.vbmf_sparse(Y_toy, params_sparse_init, 100, diag_var = false, verb = true,
    logdir = data_path, desc = "sparse_test", full_cov = true);
err = norm(Y_toy - sparse_vbmf.YHat)
println("||Y - Yhat|| = $err")
println("")
# compare the results

println(" original A          basic reconstruction          sparse reconstruction with true start")
for m in 1:M
    println(A_toy[m,:],"   ", res_vbmf.AHat[m,:], "   ", sparse_vbmf.AHat[m,:])
end


export cg

function cg(A::Union(AbstractMatrix, Function),b::Vector;
	tol::Real=1e-6,maxIter::Int64=1000,
	M::Union(AbstractMatrix, Function)=x->x ,x0::Vector=[],out::Int64=0)

	n = length(b)
	nb = norm(b)

	# allocate memory once
	Ap = zeros(eltype(b),n)

	Af =  isa(A,Function) ? A : x->A_mul_B!(Ap, A, x)
	Mf =  isa(M,Function) ? M : x->M\x

	if isempty(x0)
		x0 = zeros(eltype(b),n)
		r = copy(b)
	else
		r = b - Af(x0)
	end

	z = Mf(r)
	p = copy(z)

	resvec = zeros(maxIter)

	iter   = 1
	flag   = 1

	@inbounds for iter=1:maxIter
		Ap = Af(p)
		gamma = dot(r,z)
		alpha = gamma/dot(p,Ap)

		# if A is ill-conditioned
		if alpha==Inf || alpha<0
			flag = 2; break
		end

		BLAS.axpy!(n,alpha,p,1,x0,1) # x = alpha*p+x
		BLAS.axpy!(n,-alpha,Ap,1,r,1) # r -= alpha*Ap

		resvec[iter] = norm(r)/nb
		# converge
		if  resvec[iter] <= tol
			flag = 0; break
		end

		z    = Mf(r)
		beta = dot(z,r)/gamma

		# p = z + beta*p
		p = BLAS.scal!(n,beta,p,1)
		p = BLAS.axpy!(n,1.0,z,1,p,1)
	end
	# control output
	if out>=0
		if flag == 1
			println("Does not converge within maximum iterations.")
		elseif flag == 2
			println("Matrix A in cg has to be positive definite.")
		end
	end
	return x0,flag, resvec[iter],  resvec, iter
end


using Plots

function michalewicz(x; m=10)
    return -sum(sin(v)*sin(i*v^2/π)^(2m) for
        (i,v) in enumerate(x));
end

# plot

heatmap(0:0.1:4, 0:0.1:4, (x, y)->michalewicz([x, y]))


"""
    cubed_sphere_warp(::EquiangularSphereWarp, a, b, c)

Given points `(a, b, c)` on the surface of a cube, warp the points out to a
spherical shell of radius `R` based on the equiangular gnomonic grid proposed by
[Ronchi1996]

Returns a tuple of (x1, x2, x3), where x1, x2, x3 are the corresponding coordinates on a sphere
of radius R.

The "Cubed Sphere": A New Method for the Solution of Partial Differential Equations in Spherical Geometry
C. RONCHI,*,1 R. IACONO,* AND P. S. PAOLUCCI†
JOURNAL OF COMPUTATIONAL PHYSICS 124, 93–114 (1996)

https://github.com/CliMA/ClimateMachine.jl/blob/master/src/Numerics/Mesh/Topologies.jl
"""
function cubed_sphere_warp(
    ::EquiangularSphereWarp,
    a::FT,
    b::FT,
    c::FT,
) where {FT <: AbstractFloat}

    R = max(abs(a), abs(b), abs(c))
    function f(sR, ξ, η)
        X, Y = tan(π * ξ / 4), tan(π * η / 4)
        ζ1 = sR / sqrt(X^2 + Y^2 + 1)
        ζ2, ζ3 = X * ζ1, Y * ζ1
        ζ1, ζ2, ζ3
    end

    fdim = argmax(abs.((a, b, c)))
    if fdim == 1 && a < 0
        # (-R, *, *) : formulas for Face I from Ronchi, Iacono, Paolucci (1996)
        #              but for us face IV of the developed net of the cube
        x1, x2, x3 = f(-R, b / a, c / a)
    elseif fdim == 2 && b < 0
        # ( *,-R, *) : formulas for Face II from Ronchi, Iacono, Paolucci (1996)
        #              but for us face V of the developed net of the cube
        x2, x1, x3 = f(-R, a / b, c / b)
    elseif fdim == 1 && a > 0
        # ( R, *, *) : formulas for Face III from Ronchi, Iacono, Paolucci (1996)
        #              but for us face II of the developed net of the cube
        x1, x2, x3 = f(R, b / a, c / a)
    elseif fdim == 2 && b > 0
        # ( *, R, *) : formulas for Face IV from Ronchi, Iacono, Paolucci (1996)
        #              but for us face III of the developed net of the cube
        x2, x1, x3 = f(R, a / b, c / b)
    elseif fdim == 3 && c > 0
        # ( *, *, R) : formulas for Face V from Ronchi, Iacono, Paolucci (1996)
        #              but for us face VI of the developed net of the cube
        x3, x2, x1 = f(R, b / c, a / c)
    elseif fdim == 3 && c < 0
        # ( *, *,-R) : formulas for Face VI from Ronchi, Iacono, Paolucci (1996)
        #              but for us face I of the developed net of the cube
        x3, x2, x1 = f(-R, b / c, a / c)
    else
        error(
            "invalid case for cubed_sphere_warp(::EquiangularCubedSphere): $a, $b, $c",
        )
    end
    return x1, x2, x3
end

"""
    cubed_sphere_warp(::EquidistantSphereWarp, a, b, c)

Given points `(a, b, c)` on the surface of a cube, warp the points out to a
spherical shell of radius `R` based on the equidistant gnomonic grid outlined in
[Rancic1996] and [Nair2005]

Returns a tuple of (x1, x2, x3), where x1, x2, x3 are the corresponding coordinates on a sphere
of radius R.

The "Cubed Sphere": A New Method for the Solution of Partial Differential Equations in Spherical Geometry
C. RONCHI,*,1 R. IACONO,* AND P. S. PAOLUCCI†
JOURNAL OF COMPUTATIONAL PHYSICS 124, 93–114 (1996)

A Discontinuous Galerkin Transport Scheme on the Cubed Sphere
Ramachandran D. Nair1, Stephen J. Thomas1, and Richard D. Loft1
https://doi.org/10.1175/MWR2890.1

https://github.com/CliMA/ClimateMachine.jl/blob/master/src/Numerics/Mesh/Topologies.jl
"""
function cubed_sphere_warp(
    ::EquidistantSphereWarp,
    a::FT,
    b::FT,
    c::FT,
) where {FT <: AbstractFloat}
    R = max(abs(a), abs(b), abs(c))
    r = hypot(a, b, c)
    return R * a / r, R * b / r, R * c / r
end


## Implement the Newton's method update to estimate sqrt(x)
""" the function we are using here is y^2 - x, for which
the Newton update step is y - (y^2 - x)/(2y) """
function mysqrt(x; err = max(1.0,x)*eps(typeof(x/x)), maxiter=100)
  y = one(typeof(x/x))
  iter = 0
  while abs(y^2 - x) >= err && iter <= maxiter
    y = y - (y^2 - x)/(2*y)
    iter += 1
  end
  return y
end
@show mysqrt(500)
##
@show mysqrt(0.25)


#= udemy.com code challenge
from Prof. Mike X Cohen's course in Linear Algebra
https://www.udemy.com/course/linear-algebra-theory-and-implementation

Objective:
create 2 4x6 matrices with random numbers
calculate dot products with corresponding columns
=#

using LinearAlgebra

# define 4x6 matrices - numbers range from -10 to 10
a = rand(-10:10, 4,6)
b = rand(-10:10, 4,6)

# display matrices
println("Matrix A:")
display(a)
println("\n") # 2 newlines
println("Matrix B:")
display(b)
println("\n")

# size function generates a tuple of dimensions
# in this case (4,6)
# thus rows = 1st element of tuple (4)
# cols = 2nd element of tuple (6)
rows = size(a)[1] # unused in this example

cols = size(a)[2]

# Main for loop
for i in 1:cols
  dp = dot(a[:, i], b[:, i])
  println("Dot product of matrix column $i is $(dp)")
end


#    Problem :
#    *********
#   The generalized Rosenbrock function.
#
#   Source: problem 5 in
#   S. Nash,
#   "Newton-type minimization via the Lanczos process",
#
#   SIF input: Nick Gould, Oct 1992.
#              minor correction by Ph. Shott, Jan 1995.
#
#   classification SUR2-AN-V-0
#
#   Number of variables
#
#
# Daniel Henderson, 08/2021

f = x -> begin
	N = lastindex(x)
    return 1.0 + 100sum((x[i] - x[i-1]^2)^2 for i = 2:N) + sum((x[i] - 1.0)^2 for i = 2:N)
end

g! = (g, x) -> begin
	N = lastindex(x)
	g[1] = -400 * x[1] * (x[2] - x[1]^2)
	α    =  200 * (x[2] - x[1]^2)
	for i in 2:N-1
		g[i] =  -400 * x[i] * (x[i+1] - x[i]^2) + 2 * (x[i] - 1.0) + α
		α    =   200 * (x[i+1] - x[i]^2)
	end
	g[N] = 200 * (x[N] - x[N-1]^2) + 2*(x[N] - 1.0)
	return g
end

fg! = (g, x) -> begin
	@warn "GENROSE fg! not implmented"
	# N = lastindex(x)
	# fx = 0.0
	# g[1] = -400 * x[1] * (x[2] - x[1]^2)
	# α    =  200 * (x[2] - x[1]^2)
	# for i in 2:N-1
	# 	g[i] =  -400 * x[i] * (x[i+1] - x[i]^2) + 2 * (x[i] - 1.0) + α
	# 	α    =   200 * (x[i+1] - x[i]^2)
	# 	fx  += (x[i] - x[i-1]^2)^2 + (x[i] - 1.0)^2
	# end
	# fx  += (x[N] - x[N-1]^2)^2 + (x[N] - 1.0)^2
	# g[N] = 200 * (x[N] - x[N-1]^2) + 2*(x[N] - 1.0)
	# return fx,  g
end

init = (n::Int=500) -> begin
	x0 = [j/(n+1) for j in 1:n]
	return n, x0
end

TestSet["GENROSE"] = UncProgram("GENROSE",  f, g!, fg!, init)


# Continuation of ex 7.13, Least Squares WRM

using NumericalMethodsforEngineers
using Test

@sym begin
  LagrangePolynomial(xi_, yi_) := Module([N, sum, num, den],
    begin
      N = Length(xi)
      sum = 0
      For( i=1, i <= N, Increment(i),
        begin
          num = 1
          den = 1
          For( j=1, j <= N, Increment(j),
            begin
              If(j != i, num *= (x-xi[j]))
              If(j != i, den *= (xi[i]-xi[j]))
            end
          )
          sum += yi[i] * num/den
        end
      ),
      Return(Simplify(sum))
    end
  )
end

@sym begin
  ClearAll(xi, yi, N1, Y, Ydotdot, C11, ytilde1, ytilde2)
  xi = [0, 1//2, 1]
  yi = [0, a, 1]
  Y(x_) := LagrangePolynomial(xi, yi)
  Ydotdot(x_) = D(Y(x), x, 2)
  R(x_) := Simplify(Expand(Ydotdot(x) - 3*x - 4*Y(x)))
  R(x_) = Simplify(Expand(R(x) ./ (a => -C1/4)))
  Rdot(x_) = D(R(x), C1)
  C11 = Solve(Integrate(R(x)*Rdot(x), [x, 0, 1]), C1)[1]
  SetJ(r, ToString(Simplify(R(x))))
  SetJ(C1, ToString(C11[1][2]))
  ytilde1 = Simplify(Y(x) ./ (a => -C1/4))
  ytilde2 = Simplify(ytilde1 ./ C11)
  SetJ(y, ToString(Simplify(Expand(ytilde2))));
end

println("\n\nExample 7.17: y'' = 3x + 4y, y(0)=0, y(1)=1")
println("by 1-point least squares Weighted Residual Method")
@sym Print("\nY(x) = $(Y(x))\n")
@sym Print("Y(x) = $(Simplify(Expand(Y(x) ./ (a => -C1/4))))\n")
@sym Println("R(x) = $(R(x))\n")
@sym Print("C1 = $(Solve(Integrate(R(x)*Rdot(x), [x, 0, 1]), C1)[1][1][2])\n")
@eval ytilde_1pt_leastsquares(x) = $(Meta.parse(y))
@eval rf_1pt_leastsquares(x, C1) = $(Meta.parse(r))
println()

rf_1pt_leastsquares_1(x) = rf_1pt_leastsquares(x, Meta.parse(C1))
println()

@test r == "4 + 2C1 + x + 4C1*x - 8x^2 - 4C1*x^2"
@test y == "(1/27)*x*(-8 + 35x)"
#@test (quadgk(rf_1pt_leastsquares_1, 0, 1))[1] < 5*eps()


# ------------------------------------------------------------------------------------------
# ## Going deep: Deep neural networks
#
# So far, we've learned that if we want to classify more than two fruits, we'll need to go
# beyond using a single neuron and use *multiple* neurons to get multiple outputs. We can
# think of stacking these multiple neurons together in a single neural layer.
#
# Even so, we found that using a single neural layer was not enough to fully distinguish
# between bananas, grapes, **and** apples. To do this properly, we'll need to add more
# complexity to our model. We need not just a neural network, but a *deep neural network*.
#
# There is one step remaining to build a deep neural network. We have been saying that a
# neural network takes in data and then spits out `0` or `1` predictions that together
# declare what kind of fruit the picture is. However, what if we instead put the output of
# one neural network layer into another neural network layer?
#
# This gets pictured like this below:
#
# <img src="data/deep-neural-net.png" alt="Drawing" style="width: 500px;"/>
#
# On the left we have 3 data points in blue. Those 3 data points each get fed into 4 neurons
# in purple. Each of those 4 neurons produces a single output, but those output are each fed
# into three neurons (the second layer of purple). Each of those 3 neurons spits out a
# single value, and those values are fed as inputs into the last layer of 6 neurons. The 6
# values that those final neurons produce are the output of the neural network. This is a
# deep neural network.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Why would a deep neural network be better?
#
# This is a little perplexing when you first see it. We used neurons to train the model
# before: why would sticking the output from neurons into other neurons help us fit the data
# better? The answer can be understood by drawing pictures. Geometrically, the matrix
# multiplication inside of a layer of neurons is streching and rotating the axis that we can
# vary:
#
# [Show linear transformation of axis, with data]
#
# A nonlinear transformation, such as the sigmoid function, then adds a bump to the line:
#
# [Show the linear transformed axis with data, and then a bumped version that fits the data
# better]
#
# Now let's repeat this process. When we send the data through another layer of neurons, we
# get another rotation and another bump:
#
# [Show another rotation, then another bump]
#
# Visually, we see that if we keep doing this process we can make the axis line up with any
# data. What this means is that **if we have enough layers, then our neural network can
# approximate any model**.
#
# The trade-off is that with more layers we have more parameters, so it may be harder (i.e.
# computationally intensive) to train the neural network. But we have the guarantee that the
# model has enough freedom such that there are parameters that will give the correct output.
#
# Because this model is so flexible, the problem is reduced to that of learning: do the same
# gradient descent method on this much larger model (but more efficiently!) and we can make
# it classify our data correctly. This is the power of deep learning.
# ------------------------------------------------------------------------------------------

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/intro-to-ml/17. ML - Introduction to deep learning.jl


"""
`smooth(y, span=5)`

Smooth a vector using a moving average filter. Endpoints are handled by
collapsing the length of the filter as showed below.
```
yy[1] = y[1]
yy[2] = (y[1] + y[2] + y[3]) / 3
yy[3] = (y[1] + y[2] + y[3] + y[4] + y[5]) / 5
yy[4] = (y[2] + y[3] + y[4] + y[4] + y[6]) / 5
```

**Arguments**
- `y`: Array of values.

**Optional Arguments**
- `span`: Filter length.

**Returns**
- `yy`: Smoothed array of values.
"""
function smooth(y, span=5)

    halfspan = floor(span/2)
    n = length(y)
    yy = fill(NaN, n)

    for i = 1:n

        ispan = min(i-1, n-i, halfspan)
        ileft = convert(Int, i-ispan)
        iright = convert(Int, i+ispan)

        yy[i] = mean(y[ileft:iright])

    end

    return yy

end


function sma(ta::TimeArray, n::Int)
    tstamps = ta.timestamp[n:end]

    vals = zeros(size(ta.values,1) - (n-1), size(ta.values,2))
    for i in 1:size(ta.values,1) - (n-1)
        for j in 1:size(ta.values,2)
            vals[i,j] = mean(ta.values[i:i+(n-1),j])
        end
    end

    cname = String[]
    cols  = colnames(ta)
    for c in 1:length(cols)
        push!(cname, string(cols[c], "_sma_", n))
    end

    TimeArray(tstamps, vals, cname, ta.meta)
end

function ema(ta::TimeArray, n::Int; wilder=false)
    k = if wilder
        1 / n
    else
        2 / (n + 1)
    end

    tstamps = ta.timestamp[n:end]

    vals    =  zeros(size(ta.values,1), size(ta.values,2))
    # seed with first value with an sma value
    vals[n,:] = sma(ta, n).values[1,:]

    for i in n+1:size(ta.values,1)
        for j in 1:size(ta.values,2)
            vals[i,j] = ta.values[i,j] * k + vals[i-1, j] * (1-k)
        end
    end

    cname   = String[]
    cols    = colnames(ta)
    for c in 1:length(cols)
        push!(cname, string(cols[c], "_ema_", n))
    end

    TimeArray(tstamps, vals[n:length(ta),:], cname, ta.meta)
end

function kama(ta::TimeArray, n::Int=10, fn::Int=2, sn::Int=30)
    vola = moving(sum, abs.(ta .- lag(ta)), n)
    change = abs.(ta .- lag(ta, n))
    er = safediv.(change, vola)  # Efficiency Ratio

    # Smooth Constant
    sc = (er .* (2 / (fn + 1) - 2 / (sn + 1)) .+ 2 / (sn + 1)).^2

    cl = ta[n+1:end]
    vals = similar(Array{Float64}, indices(cl.values))
    # using simple moving average as initial kama
    pri_kama = mean(ta[1:n].values, 1)

    @assert length(cl) == length(sc)

    for idx ∈ 1:length(cl)
        vals[idx, :] =
            pri_kama =
            pri_kama .+ sc[idx].values .* (cl[idx].values .- pri_kama)
    end

    cols =
    if length(ta.colnames) == 1
        ["kama"]
    else
        ["$c\_kama" for c in ta.colnames]
    end

    TimeArray(cl.timestamp, vals, cols)
end

function env(ta::TimeArray, n::Int; e::Float64 = 0.1)
    tstamps = ta.timestamp[n:end]

    s = sma(ta, n)

    upper = s.values .* (1 + e)
    lower = s.values .* (1 - e)

    up_cname = string.(colnames(ta), "_env_$n", "_up")
    lw_cname = string.(colnames(ta), "_env_$n", "_low")

    u = TimeArray(tstamps, upper, up_cname, ta.meta)
    l = TimeArray(tstamps, lower, lw_cname, ta.meta)

    merge(l, u, :inner)
end

# Array dispatch for use by other algorithms

function sma(a::Array, n::Int)
    vals = zeros(size(a,1) - (n-1), size(a,2))

    for i in 1:size(a,1) - (n-1)
        for j in 1:size(a,2)
            vals[i,j] = mean(a[i:i+(n-1),j])
        end
    end

    vals
end

function ema(a::Array, n::Int; wilder=false)
    k = if wilder
        1 / n
    else
        2 / (n + 1)
    end

    vals = zeros(size(a,1), size(a,2))
    # seed with first value with an sma value
    vals[n,:] = sma(a, n)[1,:]

    for i in n+1:size(a,1)
        for j in 1:size(a,2)
            vals[i,j] = a[i,j] * k + vals[i-1, j] * (1-k)
        end
    end

    vals[n:end, :]
end

function env(a::AbstractArray, n::Int; e::Float64 = 0.1)
    s = sma(a, n)

    upper = @. s * (1 + e)
    lower = @. s * (1 - e)

    [lower upper]
end

doc"""
    sma(arr, n)

Simple Moving Average

```math
SMA = \frac{\sum_i^n{P_i}}{n}
```
"""
sma

doc"""
    ema(arr, n, wilder=false)

Exponemtial Moving Average

A.k.a. exponentially weighted moving average (EWMA)

```math
    \text{Let } k \text{denote the degree of weighting decrease}
```

If parameter `wilder` is `true`, ``k = \frac{1}{n}``,
else ``k = \frac{2}{n + 1}``.

```math
    EMA_t = k \times P_t + (1 - k) \times EMA_{t - 1}
```
"""
ema

doc"""

Kaufman's Adaptive Moving Average

**Arguments**:

- `n`: period

- `fn`: the fastest EMA constant

- `sn`: the slowest EMA constant

**Formula**:

```math
    \begin{align*}
        KAMA_t & = KAMA_{t-1} + SC \times (Price - KAMA_{t-1}) \\
        SC     & =
            (ER \times (\frac{2}{fn + 1} - \frac{2}{sn + 1}) + \frac{2}{sn + 1})^2 \\
        ER     & = \frac{Change}{Volatility} \\
        Change & = | Price - Price_{t-n} | \\
        Volatility & = \sum_{i}^{n} | Price_i - Price_{i-1} |
    \end{align*}
```
"""
kama

doc"""

    env(arr, n; e = 0.1)

Moving Average Envelope

```math
  \begin{align*}
    \text{Upper Envelope} & = \text{n period SMA } \times (1 + e) \\
    \text{Lower Envelope} & = \text{n period SMA } \times (1 - e)
  \end{align*}
```

**Arguments**

- `e`: the envelope, `0.1` implies the `10%` envelope.

**Reference**

- [TradingView](https://www.tradingview.com/wiki/Envelope_(ENV))
"""
env