# Implement code blocks 3.11 through 3.19 in Statistical Rethinking, 2nd edition.

using Distributions
using Statistics
using StatsBase
using LinearAlgebra

grid_size = 1000
num_observations = 3
num_samples = 1000

# Code block 3.11
param_grid = range(0, 1, length=grid_size)
prior = ones((grid_size))
likelihood = [pdf(Binomial(3, p), num_observations) for p in param_grid]
posterior = likelihood .* prior
posterior = posterior / sum(posterior)

@assert abs(sum(posterior) - 1) < 1e-4 "Posterior does not sum to 1."
@assert minimum(posterior) >= 0
@assert maximum(posterior) <= 1

samples = sample(param_grid, Weights(posterior), num_samples; replace=true, ordered=false)

# Code block 3.12
function central_interval(samples, proba::AbstractFloat)
    @assert proba > 0
    @assert proba < 1
    quantile(samples, [0.5 - proba / 2, 0.5 + proba / 2])
end

central_interval(samples, 0.5)

# Code block 3.13

# UPDATE: all the prob. prog. packages have it right. Because they *sort* the samples
# first, they're really working with the ECDF. Subtracting in the domain of the ECDF
# such that the range always has at least `proba` mass.

# - I'm 85% sure the StatisticalRethinking.jl package has this wrong, it assumes
#   some mass on both sides of the tail, which is not correct.

function high_density_interval(samples, proba::AbstractFloat)
    @assert proba >= 0
    @assert proba <= 1

    count_threshold = proba * length(samples)
    hist = fit(Histogram, samples; closed=:right, nbins=100)

    interval_mass = 0
    start_index = 1
    min_interval_length = Inf
    candidate = (NaN, NaN)
    candidate_mass = 0

    for (i, w) in enumerate(hist.weights)
        interval_mass += w

        while interval_mass >= count_threshold
            interval_length = i - start_index

            if (interval_length < min_interval_length) ||
                    ((interval_length == min_interval_length) &&
                     interval_mass > candidate_mass)

                candidate = (start_index, i)
                candidate_mass = interval_mass
                min_interval_length = interval_length
            end

            interval_mass = interval_mass - hist.weights[start_index]
            start_index += 1
        end
    end

    interval = [hist.edges[1][candidate[1]], hist.edges[1][candidate[2]]]
end

high_density_interval(samples, 0.5)

# Code block 3.14 - MAP of the explicit posterior
_, i = findmax(posterior)
posterior_mode = param_grid[i]

# Code block 3.15 - MAP based on posterior samples
# - Based on a density estimate again

# Code block 3.16
mean(samples)
median(samples)

# Code block 3.17 - mean absolute error loss for a point estimate
function mae_loss(point_estimate, param_grid, posterior)
    sum(posterior .* abs.(point_estimate .- param_grid))
end

loss = mae_loss(0.5, param_grid, posterior)

# Code blocks 3.18 and 3.19 - mean absolute error for all param values
# - The book uses `sapply` in R. I don't know how to vectorize like that yet in Julia,
#   so let's just loop.
losses = [mae_loss(p, param_grid, posterior) for p in param_grid]

_, i = findmin(losses)

best_point_estimate = param_grid[i]


"""
    kruskal_mst(g, distmx=weights(g); minimize=true)

Return a vector of edges representing the minimum (by default) spanning tree of a connected,
undirected graph `g` with optional distance matrix `distmx` using [Kruskal's algorithm](https://en.wikipedia.org/wiki/Kruskal%27s_algorithm).

### Optional Arguments
- `minimize=true`: if set to `false`, calculate the maximum spanning tree.
"""
function kruskal_mst end
# see https://github.com/mauro3/SimpleTraits.jl/issues/47#issuecomment-327880153 for syntax
@traitfn function kruskal_mst(g::AG::(!IsDirected),
    distmx::AbstractMatrix{T}=weights(g); minimize=true) where {T <: Real, U, AG <: AbstractGraph{U}}

    connected_vs = IntDisjointSets(nv(g))

    mst = Vector{edgetype(g)}()
    sizehint!(mst, nv(g) - 1)

    weights = Vector{T}()
    sizehint!(weights, ne(g))
    edge_list = collect(edges(g))
    for e in edge_list
        push!(weights, distmx[src(e), dst(e)])
    end

    weight = 0
    for e in edge_list[sortperm(weights; rev=!minimize)]
        if !in_same_set(connected_vs, src(e), dst(e))
            union!(connected_vs, src(e), dst(e))
            weight += distmx[src(e), dst(e)]
            push!(mst, e)
            (length(mst) >= nv(g) - 1) && break
        end
    end

    return (mst=mst, weight=weight)
end


# back-substitution jl

"""
Performs Rx = b by back_substitution

# Arguments
- `R::Array{Float64,2}`: diagonal matrix of size m x m
- `x::Array{Float64,1}`: vector of size m. The result is stored in x.
- `b::Array{Float64,1}`: vector of size m.

# Complexity

O(m^2)

"""
function back_substitution_v!(R, x, b)
    (_, m) = size(R)

    x .= 0

    # O(m^2)
    @inbounds for i = m:-1:1
        b_i = b[i]
        for k = i + 1:m
            b_i = b_i - R[i,k] * x[k]
        end
        x[i] = b_i / R[i,i]
    end

end

"""
Performs RW = V by back_substitution

# Arguments
- `R::Array{Float64,2}`: diagonal matrix of size m x m
- `V::Array{Float64,2}`: matrix of size m x k
- `W::Array{Float64,2}`: matrix of size m x k. The result is stored in W.

# Complexity

O(km^2)

"""
function back_substitution!(R,V,W)

    (_,k) = size(V)
    (m,_) = size(R)

    #Complexity: O(km^2)
    @views @inbounds for i = 1 : k
        V[:,i] .= 0
        for j = m:-1:1
            V[j,i] = W[j,i]/R[j,j];
            @. W[1:j-1,i] = W[1:j-1,i] - R[1:j-1,j]*V[j,i];
        end
    end

end


function gen_β(β, η, D, R)
    α = 2.0 .- β .^ (-  η - 1.0 )
    mask = R .<= 1.0 ./ α
    s = 1.0 / (η + 1.0)
    βq = [ mask[i] ?  (R[i] * α[i])^s : (1.0 / (2.0 - R[i]*α[i]))^s for i in 1:D]
    βq
end

"""
    SBX_crossover(vector1, vector2, bounds, η=15, p_variable = 0.9)

Simulated binomial crossover for given two `Vectors{Real}`.
"""
function SBX_crossover(vector1, vector2, bounds, η=15, p_variable = 0.9)
    xu = view(bounds, 2,:)
    xl = view(bounds, 1,:)
    D = length(vector1)

    do_crossover = ones(Bool, D)
    do_crossover[rand(D) .> p_variable] .= false
    do_crossover[ abs.( vector2 - vector1 ) .<= eps() ] .= false

    y1 = min.( vector1, vector2 )
    y2 = max.( vector1, vector2 )
    Δ = max.(eps(), y2 - y1)

    R = rand(D)

    β = @. 1.0 + (2.0 * (y1 - xl) / Δ)
    βq = gen_β(β, η, D, R)
    c1 = @. 0.5*(y1 + y2 -  βq*Δ)

    β = @. 1.0 + (2.0 * (y1 - xl) / Δ)
    βq = gen_β(β, η, D, R)
    c2 = @. 0.5*(y1 + y2 +  βq*Δ)

    # swap
    mask = rand(Bool, D)
    cc = copy(c1)
    c1[mask] = c2[mask]
    c2[mask] = cc[mask]

    cc1 = copy(vector1)
    cc1[do_crossover] = c1[do_crossover]
    cc2 = copy(vector2)
    cc2[do_crossover] = c2[do_crossover]

    reset_to_violated_bounds!(cc1, bounds)
    reset_to_violated_bounds!(cc2, bounds)

    return cc1, cc2
end

"""
    DE_crossover(x, u, CR)

Binomial crossover between x and u for Differential Evolution with probability CR, i.e.,
`v[j] = u[j]` if `rand() < CR`, otherwise `v[j] = x[j]`. Return `v`.
"""
function DE_crossover(x, u, CR)
    D = length(x)
    # binomial crossover
    v = zeros(D)
    j_rand = rand(1:D)

    # binomial crossover
    for j = 1:D
        if rand() < CR || j == j_rand
            v[j] = u[j]
        else
            v[j] = x[j]
        end
    end

    return v
end


# Ranking based Metrics for Recommender Systems

"""
    ranking_stats_k(y_rec, y_rel, k = 10)

Evaluates the relevancy of top k recommendations using `precison@k`, `recall@k` and `f1_score@k`. Returns result as a `Dict`.

Here, `y_rec` are predicted probabilities for recommendation and `y_rel` defines as `1` if particular result is relevant, else `0`.
The shape of `y_rec` and `y_rel` are expected to be (1, N_elements).<br>

* `precison_k` is evaluated as `Recommended_items_that_are_relevant / Total_Recommended_items`.
* `recall_l` is evaluated as `Recommended_items_that_are_relevant / Total_Relevant_items`.
* `f1_k` is evaluated as `2 * Recommended_items_that_are_relevant / (Total_Recommended_items + Total_Relevant_items)`.
"""
function ranking_stats_k(y_rec, y_rel; k = 10)
    @assert size(y_rec) == size(y_rel)
    if size(y_rec,1) !=1
       y_rec = reshape(y_rec, 1, :)
       y_rel = reshape(y_rel, 1, :)
    end
    top_k = partialsortperm(y_rec[1,:], 1:k, lt= >)
    y_rec_k, y_rel_k = y_rec[:, top_k], y_rel[:, top_k]
    tp = sum(y_rel_k .== 1)
    total_positive = sum(y_rel.==1)
    precision_k = k!=0 ? tp / k : 1
    recall_k = total_positive!=0 ? tp / total_positive : 1
    f1_k = (total_positive + k!=0) ? (2 * tp / (total_positive + k)) : 1
    return Dict("precision_k" => precision_k,
                "recall_k" => recall_k,
                "f1_k" => f1_k)
end

"""
    avg_precision(y_rec, y_rel, k = 10)

Evaluates how much of the relevant documents are concentrated in the highest ranked predictions.

Calculated as
    `∑(Recall@i - Recall@i-1)* Precision@i for i = (1, 2, 3....k)`

Here, `y_rec` are predicted probabilities for recommendation and `y_rel` defines as `1` if particular result is relevant, else `0`.
The shape of `y_rec` and `y_rel` are expected to be (1, N_elements)
"""
function avg_precision(y_rec, y_rel; k = 10)
    @assert size(y_rec) == size(y_rel)
    if size(y_rec,1) !=1
       y_rec = reshape(y_rec, 1, :)
       y_rel = reshape(y_rel, 1, :)
    end
    top_k = partialsortperm(y_rec[1,:], 1:k, lt= >)
    y_rec_k, y_rel_k = y_rec[:, top_k], y_rel[:, top_k]
    correct_prediction = 0
    running_sum = 0
    for i in 1:k
        if y_rel_k[i] == 1
            correct_prediction +=1
            running_sum += correct_prediction / i
        end
    end
    return running_sum / sum(y_rel .== 1)
end

# TODO

# DCG
# NDCG


import Distributions: cdf, Normal

N(x) = cdf(Normal(), x)

function d1(S, K, τ, r, σ, q)
    return (log(S / K) + (r - q + σ^2 / 2) * τ) / (σ * √(τ))
end

function d2(S, K, τ, r, σ, q)
    return d1(S, K, τ, r, σ, q) - σ * √(τ)
end

"""
    eurocall(;S=1.,K=1.,τ=1,r,σ,q=0.)

Calculate the Black-Scholes implied option price for a european call, where:

- `S` is the current asset price
- `K` is the strike or exercise price
- `τ` is the time remaining to maturity (can be typed with \\tau[tab])
- `r` is the continuously compounded risk free rate
- `σ` is the (implied) volatility (can be typed with \\sigma[tab])
- `q` is the continuously paid dividend rate

Rates should be input as rates (not percentages), e.g.: `0.05` instead of `5` for a rate of five percent.

!!! Experimental: this function is well-tested, but the derivatives functionality (API) may change in a future version of ActuaryUtilities.

# Extended Help

This is the same as the formulation presented in the [dividend extension of the BS model in Wikipedia](https://en.wikipedia.org/wiki/Black%E2%80%93Scholes_model#Black%E2%80%93Scholes_equation).

## Other general comments:

- Swap/OIS curves are generally better sources for `r` than government debt (e.g. US Treasury) due to the collateralized nature of swap instruments.
- (Implied) volatility is characterized by a curve that is a function of the strike price (among other things), so take care when using
- Yields.jl can assist with converting rates to continuously compounded if you need to perform conversions.

"""
function eurocall(; S = 1.0, K = 1.0, τ = 1, r, σ, q = 0.0)
    d₁ = d1(S, K, τ, r, σ, q)
    d₂ = d2(S, K, τ, r, σ, q)
    return (N(d₁) * S * exp(τ * (r - q)) - N(d₂) * K) * exp(-r * τ)
end

"""
    europut(;S=1.,K=1.,τ=1,r,σ,q=0.)

Calculate the Black-Scholes implied option price for a european call, where:

- `S` is the current asset price
- `K` is the strike or exercise price
- `τ` is the time remaining to maturity (can be typed with \\tau[tab])
- `r` is the continuously compounded risk free rate
- `σ` is the (implied) volatility (can be typed with \\sigma[tab])
- `q` is the continuously paid dividend rate

Rates should be input as rates (not percentages), e.g.: `0.05` instead of `5` for a rate of five percent.

!!! Experimental: this function is well-tested, but the derivatives functionality (API) may change in a future version of ActuaryUtilities.

# Extended Help

This is the same as the formulation presented in the [dividend extension of the BS model in Wikipedia](https://en.wikipedia.org/wiki/Black%E2%80%93Scholes_model#Black%E2%80%93Scholes_equation).

## Other general comments:

- Swap/OIS curves are generally better sources for `r` than government debt (e.g. US Treasury) due to the collateralized nature of swap instruments.
- (Implied) volatility is characterized by a curve that is a function of the strike price (among other things), so take care when using
- Yields.jl can assist with converting rates to continuously compounded if you need to perform conversions.

"""
function europut(; S = 1.0, K = 1.0, τ = 1, r, σ, q = 0.0)
    d₁ = d1(S, K, τ, r, σ, q)
    d₂ = d2(S, K, τ, r, σ, q)
    return (N(-d₂) * K - N(-d₁) * S * exp(τ * (r - q))) * exp(-r * τ)
end


function KalmanFilter(T::Matrix, Z::Matrix, Q::Matrix, H::Matrix, n::Integer, α0, y)

    # Initialize state variables
    αhat = repeat([zeros(size(α0))], n);
    αhat[1] = T*α0;

    println("I'm here 1")

    yhat = repeat([zeros(size(y[1]))], n);
    v = copy(yhat);

    P0 = [2.0]; # Initial guess given in the problem

    P = repeat([zeros(size(Q))], n);
    P[1] = T * P0 * T' + Q;

    println("I'm here 2, P[1] = $(P[1])")

    F = repeat([zeros(size(H))], n)
    G = repeat([zeros(size(Z * Q))], n)
    K = repeat([zeros(size(T * G[1]' * ones(size(H))))], n)

    for i in 1:n-1
        yhat[i] = Z * αhat[i];
        v[i] = y[i] - yhat[i];

        F[i] = Z * P[i] * Z' + H;
        G[i] = Z * P[i];
        K[i] = T * G[i]' * inv(F[i]);

        αhat[i + 1] = T * αhat[i] + K[i] * v[i];
        P[i + 1] = T * (P[i] - G[i]' * inv(F[i]) * G[i]) * T' + Q;
    end

    yhat[n] = Z * αhat[n];
    v[n] = y[n] - yhat[n];

    F[n] = Z * P[n] * Z' + H;
    G[n] = Z * P[n];
    K[n] = T * G[n]' * inv(F[n]);

    return αhat, v, yhat, P, F, G, K
end


###########################################
# Unconstrained optimization
###########################################

f(x) = sin(x[1] + x[2]) + cos(x[1])^2
g(x) = [cos(x[1] + x[2]) - 2*cos(x[1])*sin(x[1]); cos(x[1] + x[2])]

f(x1,x2) = f([x1;x2])

f([0; 0])
f(0, 0)

#### Exercise

using Plots

xs = range(-3, 1, length = 40)
ys = range(-2, 1, length = 40)

contourf(xs, ys, f, color = :jet)

#### Exercise

finite_difference(f, x::Real; h=1e-8) = (f(x+h) - f(x)) / h

#### Exercise

x = [-2; -1]
fin_diff(h) = finite_difference(y -> f(x[1], y), x[2]; h=h)

true_grad = g(x)[2]

hs = 10. .^ (-15:0.01:-1)

plot(hs, fin_diff,
    xlabel = "h",
    ylabel = "Partial gradient wrt y",
    label = ["Approximation" "True gradient"],
    xscale = :log10,
)

hline!([true_grad]; label =  "True gradient")

#### Numerical errors

x = 1
h = 1e-13

(x+h)^2 - x^2
2*x*h + h^2

#### Exercise

x = [-2; -1]
α = 0.25
x_grad = [x x.+α.*g(x)]

contourf(xs, ys, f; color = :jet)
plot!(x_grad[1, :], x_grad[2, :];
    line = (:arrow, 4, :black),
    label = "",
)

#### Exercise

function optim(f, g, x, α; max_iter=100)
    xs = zeros(length(x), max_iter+1)
    xs[:,1] = x
    for i in 1:max_iter
        x -= α*g(x)
        xs[:,i+1] = x
    end
    return xs
end

#### Animation

using Random

function create_anim(
    f,
    path,
    xlims,
    ylims,
    file_name = joinpath(pwd(), randstring(12) * ".gif");
    xbounds = xlims,
    ybounds = ylims,
    fps = 15,
)
    xs = range(xlims...; length = 100)
    ys = range(ylims...; length = 100)
    plt = contourf(xs, ys, f; color = :jet)

    # add constraints if provided
    if !(xbounds == xlims && ybounds == ylims)
        x_rect = [xbounds[1]; xbounds[2]; xbounds[2]; xbounds[1]; xbounds[1]]
        y_rect = [ybounds[1]; ybounds[1]; ybounds[2]; ybounds[2]; ybounds[1]]

        plot!(x_rect, y_rect; line = (2, :dash, :red), label="")
    end

    # add an empty plot
    plot!(Float64[], Float64[]; line = (4, :arrow, :black), label = "")

    # extract the last plot series
    plt_path = plt.series_list[end]

    # create the animation and save it
    anim = Animation()
    for x in eachcol(path)
        push!(plt_path, x[1], x[2]) # add a new point
        frame(anim)
    end
    gif(anim, file_name; fps = fps, show_msg = false)
    return nothing
end

#### Exercise

x_gd = optim([], g, [0; -1], 0.1)

xlims = (-3, 1)
ylims = (-2, 1)
create_anim(f, x_gd, xlims, ylims, "anim1.gif")

f_gd = [f(x) for x in eachcol(x_gd)]

plot(f_gd, label="", xlabel="Iteration", ylabel="Function value")

#### Different stepsizes

x_gd = optim([], g, [0; -1], 0.01)

create_anim(f, x_gd, xlims, ylims, "anim2.gif")

x_gd = optim([], g, [0; -1], 1)

create_anim(f, x_gd, xlims, ylims, "anim3.gif")

#### Stepsize selection

abstract type Step end

struct GD <: Step
    α::Float64
end

#### Gradient descent

optim_step(s::GD, f, g, x) = -s.α*g(x)

function optim(f, g, x, s::Step; max_iter=100)
    xs = zeros(length(x), max_iter+1)
    xs[:,1] = x
    for i in 1:max_iter
        x += optim_step(s, f, g, x)
        xs[:,i+1] = x
    end
    return xs
end

gd = GD(0.1)
x_opt = optim(f, g, [0;-1], gd)

create_anim(f, x_opt, xlims, ylims, "anim4.gif")

#### Exercise

struct Armijo <: Step
    c::Float64
    α_max::Float64
end

function optim_step(s::Armijo, f, g, x)
    fun = f(x)
    grad = g(x)
    α = s.α_max
    while f(x .- α*grad) > fun - s.c*α*(grad'*grad)
        α /= 2
        if α <= 1e-6
            warning("Armijo line search failed.")
            break
        end
    end
    return -α*grad
end

gd = Armijo(1e-4, 1)
x_opt = optim(f, g, [0;-1], gd)

create_anim(f, x_opt, xlims, ylims, "anim5.gif")

###########################################
# Constrained optimization
###########################################

f(x) = sin(x[1] + x[2]) + cos(x[1])^2
g(x) = [cos(x[1] + x[2]) - 2*cos(x[1])*sin(x[1]); cos(x[1] + x[2])]

f(x1,x2) = f([x1;x2])

#### Projected gradients

function optim(f, g, P, x, α; max_iter=100)
    xs = zeros(length(x), max_iter+1)
    ys = zeros(length(x), max_iter)
    xs[:,1] = x
    for i in 1:max_iter
        ys[:,i] = xs[:,i] - α*g(xs[:,i])
        xs[:,i+1] = P(ys[:,i])
    end
    return xs, ys
end

P(x, x_min, x_max) = min.(max.(x, x_min), x_max)

x_min = [-1; -1]
x_max = [0; 0]

xs, ys = optim(f, g, x -> P(x,x_min,x_max), [0;-1], 0.1)

#### Plot 1

xlims = (-3, 1)
ylims = (-2, 1)

create_anim(f, xs, xlims, ylims, "anim6.gif";
    xbounds=(x_min[1], x_max[1]),
    ybounds=(x_min[2], x_max[2]),
)

#### Plot 2

xys = hcat(reshape([xs[:,1:end-1]; ys][:], 2, :), xs[:,end])

create_anim(f, xys, xlims, ylims, "anim7.gif";
    xbounds=(x_min[1], x_max[1]),
    ybounds=(x_min[2], x_max[2]),
)


@doc raw"""
    gauss_quadrature(q)

Return the Gauss-Legendre quadrature rule with `q` points (integrates polynomials of degree
$2q-1$ exactly).

A tuple `(x,w)` is returned.
"""
function gauss_quadrature(q)
    x = zeros(CeedScalar, q)
    w = zeros(CeedScalar, q)
    C.CeedGaussQuadrature(q, x, w)
    x, w
end

struct QuadratureMode{T} end
const Abscissa = QuadratureMode{:Abscissa}()
const AbscissaAndWeights = QuadratureMode{:AbscissaAndWeights}()

@doc raw"""
    lobatto_quadrature(q, mode::Mode=Abscissa)

Return the Gauss-Lobatto quadrature rule with `q` points (integrates polynomials of degree
$2q-3$ exactly).

If `mode` is `AbscissaAndWeights`, then both the weights and abscissa are returned as a
tuple `(x,w)`.

Otherwise, (if `mode` is `Abscissa`), then only the abscissa `x` are returned.
"""
function lobatto_quadrature(q, mode::Mode=Abscissa) where {Mode}
    return_weights = (mode == AbscissaAndWeights)
    x = zeros(CeedScalar, q)
    w = (return_weights) ? zeros(CeedScalar, q) : C_NULL
    C.CeedLobattoQuadrature(q, x, w)
    return_weights ? (x, w) : x
end


# ---
# title: 50. Pow(x, n)
# id: problem50
# author: zhwang
# date: 2022-03-02
# difficulty: Medium
# categories: Math, Binary Search
# link: <https://leetcode.com/problems/powx-n/description/>
# hidden: true
# ---
#
# Implement [pow( _x_ , _n_
# )](http://www.cplusplus.com/reference/valarray/pow/), which calculates  _x_
# raised to the power _n_ (i.e. xn).
#
#
#
# **Example 1:**
#
#
#
#     Input: x = 2.00000, n = 10
#     Output: 1024.00000
#
#
# **Example 2:**
#
#
#
#     Input: x = 2.10000, n = 3
#     Output: 9.26100
#
#
# **Example 3:**
#
#
#
#     Input: x = 2.00000, n = -2
#     Output: 0.25000
#     Explanation: 2-2 = 1/22 = 1/4 = 0.25
#
#
#
#
# **Constraints:**
#
#   * `-100.0 < x < 100.0`
#   * `-231 <= n <= 231-1`
#   * `-104 <= xn <= 104`
#
#
## @lc code=start
using LeetCode

function my_pow(x::Float64, n::Int)::Float64
    n == 0 && return 1
    if n < 0
        x, n = 1 / x, -n
    end
    res, extra = x, 1
    while n > 1
        n & 1 == 1 && (extra *= res)
        res *= res
        n >>= 1
    end
    return res * extra
end
## @lc code=end


function EuropeanDownAndOutCallMC(K, T, S, sig, r, div, H, N, M)

dt = T/N
nudt = (r - div - 0.5 * sig^2) * dt
sigsdt = sig * sqrt(dt)

## variables to store the sums of the option values

sum_CT = 0
sum_CT2 = 0

## outer loop is the number of simulations bigger is better

for j = 1:M

        St = S
        BARRIERCROSSED = false
    ## inner loop is simulating each price path of length N
        for i = 1:N

        St = St * exp( nudt + sigsdt * randn()) ## note, it only add the changes, does not store the path

        if St <= H

                BARRIERCROSSED = true
                break
        end

        end

    ## wraps up and stores all the values of each path

if BARRIERCROSSED == true
            CT = 0

        else

            CT = max(0, St - K)
        end

sum_CT = sum_CT + CT
sum_CT2 = sum_CT2 + CT^2

    end

  ## discounting the average call value
call_value = sum_CT / M * exp(-r*T)

SD = (sqrt(sum_CT2 - 1/M * (sum_CT^2)) * exp(-2*r*T)) / (M-1)
SE = SD / sqrt(M)

println("The Call Value is $call_value")
println("The SD is $SD")
println("The SE is $SE")

MC = [call_value, SD, SE]

end


function MultBinTreeAMPut(K, T, S, r, N, u, d)
        dt = T/N ## Time to Maturity/ Number of Steps
        p = (exp(r * dt) - d) / (u - d)
        disc = exp(-r * dt)

        #St = linspace(1,N+1,N+1)

        St = Array(Float64, (N+1, N+1))

        # Initialising asset prices at maturity

        St[1 , N+1] = S * u^N
        for j in 2:(N+1)
                St[j, N+1] = St[j-1, N+1] * d/u
        end

        #St = flipdim(St, 1)

        # create Array for options value tree

        P = Array(Float64, (N+1, N+1))#linspace(1, N+1, N+1)

        # Initialise Option Values at maturity

        for j in 1:(N+1)
                P[j, N+1] = max(0.0, K - St[j, N+1])
        end

        # Stepping back through the tree

        for i = linrange(N, 1, N)
                for j in 1:i
                      P[j, i] = disc * (p * P[j, i+1] + (1-p) * P[j+1, i+1])
                      St[j,i] = St[j,i+1] * d
                      P[j, i] = max(P[j, i], K - St[j,i])
                end
        end

P

end

MultBinTreeAMPut(100,1,100,0.06,3,1.1,0.9091)


export factors
export primefactors
export totient

"""
Return sorted list of unique factors for a number n.
"""
function factors(n::Integer)::Array{Integer}
    factors = Set{Integer}(1)
    for i in 2:Integer(floor(sqrt(n)))
        if n % i == 0
            push!(factors, i)
            push!(factors, div(n, i))
        end
    end
    return sort(collect(factors))
end

"""
Return sorted list of prime factors for number n.
"""
function primefactors(n::Integer)::Array{Integer}
    factors = ones(Integer, 0)

    # Store all of the factors of 2.
    while n % 2 == 0
        push!(factors, 2)
        n = fld(n, 2)
    end

    # Store all other prime factors, increasing to sqrt(n).
    for i in 3:2:Integer(floor(sqrt(n)))
        while n % i == 0
            push!(factors, i)
            n = fld(n, i)
        end
    end

    # Store last prime factor.
    if n > 1
        push!(factors, n)
    end

    return sort(factors)
end

"""
Euler's totient function returns the number of positive integers up to n that
are releatively prime to n.  Pull out the prime factors f and take a produce of
n and all (1 - 1/f).
"""
function totient(n::Integer)::Integer
    factors = primefactors(n)
    last_prime = 0
    result = Float64(n)
    for f in factors
        if f == last_prime
            continue
        else
            result *= (1 - 1 / f)
            last_prime = f
        end
    end
    return Integer(round(result))
end


half(L) = div(L - 1, 2)

logZ(M::Matrix) = logdet(I + M)
Z(M::Matrix) = det(I + M)

function pairing(L, h)
    E(k) = sqrt(1 + h^2 + 2h * cospi((2k + 1) / L))
    ζ(k) = - (h + cospi((2k + 1) / L))
    u(k) = sqrt((1 + ζ(k) / E(k)) / 2)
    v(k) = sqrt((1 - ζ(k) / E(k)) / 2)
    g(r) = 2 / L * sum(sinpi((2k + 1) / L * r) * v(k) / u(k) for k in 0:div(L - 1, 2))

    M = zeros(L, L)
    for i in 1:L, j in 1:L
        M[i,j] = g(i -  j)
    end
    return M
end

function get_prob(x::BitVector, M::Matrix)
    amp = log(abs.(det(M[x .> 0, x .> 0])))
    return amp
end

function get_all_probs(M::Matrix)
    L = size(M, 1)
    logZ = log(abs(det(M)))
    x = i->bitarray(i, L)
    amps = [get_prob(x(i), M) for i in 0:2^L - 1]
    amps = exp.(amps)
    amps2 = normalize(amps, 1)
    return round.(amps, digits = 7), round.(amps2, digits = 7)
end

# for L in [2, 4,8,16,20, 32]
#     P = pairing(L, 1.0)
#     Z1 = logdet(I + P)
#     Z2 = log(sum(get_all_probs(P)[1]))
#     println("L is $L Z is $Z1 and $Z2")
# end


#=
Project:      MIRTdev
File:         mri_geom_compress.jl
Description:  MRI Coil Compression
Author:       Rajas Gupta
Date Created: October 8th 2019
Date Update:  October 8th 2019
=#

export ir_mri_coil_compress

using LinearAlgebra: svd, norm

"""
    (odata, σ, Vr) = ir_mri_coil_compress(idata ; ncoil)

MRI coil compression via PCA.
Given multiple MRI surface coil images (idata), use SVD/PCA
to find a smaller number of virtual coil images (odata).

In:
* `idata` `[(N) n_in]`: noisy complex images (2D or 3D) for each coil

Option:
* `ncoil` Desired # of virtual coils (default: 1)

Out:
* `odata` `[(N) ncoil]`: virtual coil images
* `σ`     `[n_in]`: singular values.
* `Vr`    `[n_in, ncoil]`: compression matrix for reducing other data.

todo: currently ignores noise correlations

"""
function ir_mri_coil_compress(idata::AbstractArray{<:Number} ; ncoil::Int = 1)

    idim = size(idata)
    n_in = idim[end]
    idata = reshape(idata, :, n_in)

    (_, σ, V) = svd(idata)

    Vr = V[:, 1:ncoil]
    odata = idata * Vr
    odata = reshape(odata, idim[1:end-1]..., ncoil)
    return (odata, σ, Vr)
end

"""
    snr2sigma(db, yb)
convert SNR in dB to noise σ for complex gaussian noise
"""
function snr2sigma(db, yb::AbstractArray{<:Complex})
    10^(-db/20) * norm(yb) / sqrt(length(yb)) / sqrt(2)
end


const ykow = [1.957e-1, 1.947e-1, 1.735e-1, 1.600e-1, 8.44e-2, 6.27e-2,
              4.56e-2,  3.42e-2, 3.23e-2, 2.35e-2, 2.46e-2]

const ukow = [4.0e0, 2.0e0, 1.0e0, 5.0e-1, 2.5e-1, 1.67e-1,
              1.25e-1, 1.0e-1, 8.33e-2, 7.14e-2, 6.25e-2]

const kowalik_osborne = let res_init=zeros(11), jac_init=zeros(11,4), x_init=[.25, .39, .415, .39]

    function res(x, r)
        for i = 1:11
            r[i] = ykow[i] - x[1]*(ukow[i] ^ 2 + x[2] * ukow[i]) / (ukow[i] ^ 2 + x[3]*ukow[i] + x[4])
        end
        return r
    end

    function jac(x, j)
        for i = 1:11
            t = -1.e0 / (ukow[i]^2 + x[3] * ukow[i] + x[4])
            j[i, 1] = t * (ukow[i]^2 + x[2] * ukow[i])
            j[i, 2] = x[1] * ukow[i]*t
            t = t * j[i, 1] * x[1]
            j[i, 3] = ukow[i] * t
            j[i, 4] = t
        end
        return j
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("kowalik_osborne", res, jac,
                                                            res_init, jac_init, x_init;
                                                            scale=scale, verbose=verbose,
                                                            print_steps=print_steps)
end


using Statistics

export testchaos01

"""
    testchaos01(φ::Vector [, cs, N0]) -> chaotic?
Perform the so called "0-1" test for chaos introduced by Gottwald and
Melbourne [1] on the timeseries `φ`.
Return `true` if `φ` is chaotic, `false` otherwise.

## Description
This method tests if the given timeseries is chaotic or not by transforming
it into a two-dimensional diffusive process. If the timeseries is chaotic,
the mean square displacement of the process grows as `sqrt(length(φ))`,
while it stays constant if the timeseries is regular.
The implementation here computes `K`, the correlation coefficient (median
of `Kc for c ∈ cs`), and simply checks if `K > 0.5`.

If you want to access the various `Kc` you should call the method
`testchaos01(φ, c::Real, N0)` which returns `Kc`.

`cs` defaults to `3π/5*rand(10) + π/4` and `N0`, the length
of the two-dimensional process, is `N0 = length(φ)/10`.

Notice that for data sampled from continous dynamical systems, some
care must be taken regarding the values of `cs`, see [1].

## References

[1] : Gottwald & Melbourne, “The 0-1 test for chaos: A review”
[Lect. Notes Phys., vol. 915, pp. 221–247, 2016.](www.doi.org/10.1007/978-3-662-48410-4_7)
"""
function testchaos01(φ::Vector, cs = 3π/5*rand(10) .+ π/4, N0 = Int(length(φ)÷10))
    K = median(testchaos01(φ, c, N0) for c in cs)
    return K > 0.5
end

function testchaos01(φ::Vector, c::Real, N0 = Int(length(φ)÷10))
    N, E = length(φ), mean(φ)
    @assert N0 ≤ N/10
    pc, qc = trigonometric_decomposition(φ, c)
    Dc = mmsd(E, pc, qc, N0, c)
    Kc = cor(Dc, 1:N0)
    return Kc
end

function trigonometric_decomposition(φ, c)
    X = promote_type(eltype(φ), typeof(c))
    N = length(φ)
    pc, qc = zeros(X, N), zeros(X, N)
    @inbounds for n in 1:N-1
        si, co = sincos(c*n)
        pc[n+1] = pc[n] + φ[n]*co
        qc[n+1] = qc[n] + φ[n]*si
    end
    return pc, qc
end

"modified mean square displacement"
function mmsd(mf, pc::Vector{T}, qc, N0, c) where T
    N, Dc = length(pc), zeros(T, N0)
    f = mf^2/(1-cos(c)) # constant factor for Dc
    for n in 1:N0
        Mc = sum((pc[n+j] - pc[j])^2 + (qc[n+j] - qc[j])^2 for j in 1:N-n)/(N-n)
        Dc[n] = Mc - f*(1 - cos(n*c))
    end
    return Dc
end


# Implement code blocks 3.2 - 3.10 in McElreath's Statistical Rethinking, 2nd edition.

using Distributions
using Plots
using Statistics
using StatsBase

grid_size = 1000
num_samples = 1000

## Code block 3.2
param_grid = range(0, 1, length=grid_size)
prior = ones((grid_size))
likelihood = [pdf(Binomial(9, p), 6) for p in param_grid]
posterior = likelihood .* prior
posterior = posterior / sum(posterior)

@assert abs(sum(posterior) - 1) < 1e-4 "Posterior does not sum to 1."
@assert minimum(posterior) >= 0
@assert maximum(posterior) <= 1

## Code block 3.3
samples = sample(param_grid, Weights(posterior), num_samples; replace=true, ordered=false)

## Code block 3.4
display(scatter(samples, xlabel="Sample number", ylabel="Proportion water (p)"))

## Code block 3.5
# use a histogram, because I don't yet know a comparable tool for density estimation in
# Julia.
display(histogram(samples, bins=:40))

## Code block 3.6
sum(posterior[param_grid .< 0.5])

## Code block 3.7
sum(samples .< 0.5) / num_samples

## Code block 3.8
sum(0.5 .< samples .< 0.75) / num_samples

## Code block 3.9
quantile(samples, 0.8)

## Code block 3.10
quantile(samples, [0.1, 0.9])


"""
    fresnels(x)

Compute the Fresnel S integral of type 0 (integral of `sin(πt²/2)`).
"""
fresnels(x) = (1 + im) / 4 * (erf((1 + im)√π / 2 * x) - im * erf((1 - im)√π / 2 * x))

"""
    fresnelc(x)

Compute the Fresnel C integral of type 0 (integral of `cos(πt²/2)`).
"""
fresnelc(x) = (1 - im) / 4 * (erf((1 + im)√π / 2 * x) + im * erf((1 - im)√π / 2 * x))


using JuMP
using Cbc
# cria o modelo para a otimização linear, utiliza o cbc como solver
m = Model(with_optimizer(Cbc.Optimizer))

set_time_limit_sec(m , 1)

n = 6

# define as variaveis binárias xij e o u como inteiro
@variable(m , x[1:n , 1:n] , Bin)
@variable(m , u[1:n] , integer=true)

fix(u[1] , 0)

# cria o grafo de pesos para fazer conexão
y = zeros(Int64 , n , n)

# dados do exemplo do slide da prof(trabalho -  material auxiliar)
# é o exemplo das cidades que está logo depois do MTZ
y = [ 0  186 105 208 177 94 ;
     187 0 89.8 223 255 254 ;
     99.9 89 0  203 204 172 ;
     206 220 203 0 377 295 ;
     168 251 201 376 0 156 ;
     86.8 255 173  293 159  0]

#for i in 1:5
#    for j in 1:5
#        y[i,j] = i+j
#    end
#end

# função objetiva
@objective(m , Min , sum(x.*y))

# constraint de que a soma de uma linha vai ser 1
@constraint(m , [i= 1:n] , sum(x[i,1:end .!=i]) == 1)
# constraint de que a soma de uma coluna vai ser 1
@constraint(m , [i= 1:n] , sum(x[1:end.!=i,i]) == 1)

# constraint do u
@constraint(m , [i = 1:n , j =2:n , i != j] , u[j] >= u[i] + x[i,j] - n *(1-x[i,j]))

#printa os dados do modelo
println(m)

# resolve e printa os resultados
optimize!(m)

uAux = value.(u)
xAux = value.(x)

for i in 1:n
    println("u$(i) $(uAux[i])")
end

print("\n\n Xij:\n")

for i in 1:n
    for j in 1:n
        print("$(xAux[i,j]) ")
    end
    println()
end


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


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)


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


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


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


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